{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Simple and Effective Gradient Based Tuning of Sequence to Sequence Models\n",
        "\n",
        "This interactive notebook accompanies the AutoML paper (https://openreview.net/pdf?id=RBTUKLfQ_pc).\n",
        "\n",
        "[Open in Google  Colab](https://colab.research.google.com/github/google-research/google-research/blob/master/speech_embedding/speech_commands.ipynb)\n",
        "\n",
        "Before running any cells please enable GPU for this notebook, otherwise training will be very slow. \n",
        "\n",
        "* *Edit* → *Notebook Settings*\n",
        "* select *GPU* from the *Hardware Accelerator* drop-down\n",
        "\n",
        "Usage:\n",
        "After connecting to a runtime, run all cells via *Runtime*→*Run all*. This will run initially just the 'Setup' cell block. This loads the required packages at the appropriate versions, and clones the code from the github. Once Setup is done, the runtime will be restarted to load the required package versions. Afterwards, again run *Runtime*→*Run all*, and this time 'Setup' will be passed through without making any changes, and the rest of the colab will be executed."
      ],
      "metadata": {
        "id": "-0esjq3luJYi"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JndnmDMp66FL"
      },
      "source": [
        "##### Copyright 2022 Google LLC.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hMqWDc_m6rUC"
      },
      "outputs": [],
      "source": [
        "#@title\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dbh3a1h0zmJ4"
      },
      "source": [
        "# Setup\n",
        "\n",
        "When running the Setup section the first time, we will automatically restart the runtime to load the required package versions. Afterwards, we will pass through without loading again or restarting.\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title\n",
        "import os\n",
        "already_cloned = bool(os.path.exists(f'{os.getcwd()}/google-research'))\n",
        "if not already_cloned:\n",
        "  # Clone the codebase. Creates 'content/google-research/' and populates with code.\n",
        "  !git clone https://github.com/google-research/google-research.git --depth=1\n",
        "  !pip3 install -r google-research/gradient_based_tuning/requirements.txt\n",
        "\n",
        "print(f'The code from github.com/google-research/google-research/* has been cloned to {os.getcwd()}/google-research.')"
      ],
      "metadata": {
        "id": "Z2kAd5EiJDYR",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title\n",
        "# Import a compatible version of jaxlib which supports running on GPU.\n",
        "\n",
        "import jax\n",
        "if jax.default_backend() != 'gpu':\n",
        "  !pip3 install https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.1.75+cuda11.cudnn82-cp37-none-manylinux2010_x86_64.whl\n",
        "  print('We have installed a new version of jaxlib that allows compilation for GPU, and will now restart the runtime to load that version.')\n",
        "  print('-----------------------')\n",
        "  print('Colab will restart now. There will be a popup in the bottom left saying the runtime has crashed unexpectedly, which can be ignored. After that happens, run all cells again.')\n",
        "  print('-----------------------')\n",
        "  import os\n",
        "  os.kill(os.getpid(), 9)\n",
        "\n",
        "print(f'JAX is running on {jax.default_backend()}.')"
      ],
      "metadata": {
        "id": "LXty3cOiBS1a",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Data"
      ],
      "metadata": {
        "id": "3O4nT_zSuUJ4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Imports\n",
        "# This is necessary to allow python imports to work properly.\n",
        "import sys\n",
        "if './google-research' not in sys.path:\n",
        "  sys.path.append('./google-research')\n",
        "\n",
        "# Load the data\n",
        "import csv\n",
        "import os\n",
        "import tempfile\n",
        "import sys\n",
        "\n",
        "from absl import app\n",
        "from absl import flags\n",
        "from sentencepiece import SentencePieceTrainer\n",
        "import tensorflow.compat.v2 as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "\n",
        "# Produce training examples\n",
        "import difflib\n",
        "\n",
        "from typing import Dict, Iterator, List, Tuple\n",
        "\n",
        "from absl import logging\n",
        "import apache_beam as beam\n",
        "from apache_beam import metrics\n",
        "from apache_beam.io import tfrecordio\n",
        "from apache_beam.options import pipeline_options\n",
        "import sentencepiece as spm\n",
        "from tensor2tensor.data_generators.generator_utils import pack_examples\n",
        "from tensor2tensor.data_generators.generator_utils import to_example\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "z6kDrYJOBmJu",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title\n",
        "# This helps to avoid OOM when running on GPU.\n",
        "# https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html#common-causes-of-oom-failures\n",
        "tf.config.experimental.set_visible_devices([], 'GPU')"
      ],
      "metadata": {
        "id": "rlXVpreL7JTu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load raw data\n",
        "\n",
        "Load raw examples from tfds. See https://www.tensorflow.org/datasets/catalog/wmt19_translate for more details. Takes about 10 min with lithuanian-english (lt-en) dataset. Will take longer with larger datasets."
      ],
      "metadata": {
        "id": "aeTXmggZ87yx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "SOURCE_LANGUAGE='lt'\n",
        "TARGET_LANGUAGE='en'\n",
        "LANG_PAIR=SOURCE_LANGUAGE+TARGET_LANGUAGE\n",
        "OUTPUT_DIR='/content/data'\n",
        "\n",
        "\n",
        "TFDS_NAME='wmt19_translate/%s-%s' % (SOURCE_LANGUAGE, TARGET_LANGUAGE)  # Which tfds dataset to use\n",
        "SPLIT_FOR_GUIDANCE_DATA=0.01\n",
        "\n",
        "LOGGING_STEPS = 100000"
      ],
      "metadata": {
        "id": "PBgJ4iA4-Kyx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title\n",
        "def _read_string_tensor(tensor):\n",
        "  raw_text = tensor.numpy().decode('utf-8')\n",
        "  return ' '.join(raw_text.strip().split())\n",
        "\n",
        "\n",
        "def write_data_to_tsv(\n",
        "    output_dir,\n",
        "    source_language,\n",
        "    target_language,\n",
        "    tfds_name,\n",
        "    split_for_guidance_data,\n",
        "):\n",
        "  \"\"\"Download data and write it to plain tsv for train, dev, and guide datasets.\n",
        "\n",
        "  Args:\n",
        "    output_dir: The dir to which the data will be written. Dirs will be\n",
        "      recursively created if not already present.\n",
        "    source_language: Source language of the translation data.\n",
        "    target_language: Target language of the translation data.\n",
        "    tfds_name: The name of the desired dataset in tfds. (ie.\n",
        "      wmt19_translate/de-en). See\n",
        "      https://www.tensorflow.org/datasets/catalog/wmt19_translate for more\n",
        "        details.\n",
        "    split_for_guidance_data: How much of the training data to set aside for\n",
        "      guidance dataset. Defaults to 1%. Set to 0 to produce a full training\n",
        "      split.\n",
        "  \"\"\"\n",
        "  lang_pair = source_language + target_language\n",
        "\n",
        "  if not 0 < split_for_guidance_data < 1:\n",
        "    raise ValueError('split_for_guidance_data must be between 0 and 1: (%s)' %\n",
        "                     split_for_guidance_data)\n",
        "  train_data_path_tsv = os.path.join(\n",
        "      output_dir, '{}_train_{:.0f}percent.tsv'.format(\n",
        "          lang_pair, 100 * (1.0 - split_for_guidance_data)))\n",
        "  guide_data_path_tsv = os.path.join(\n",
        "      output_dir,\n",
        "      '{}_guide_{:.0f}percent.tsv'.format(lang_pair,\n",
        "                                          100 * (split_for_guidance_data)))\n",
        "\n",
        "  tf.io.gfile.makedirs(output_dir)\n",
        "  guide_example_count = 0\n",
        "  train_example_count = 0\n",
        "  print('Writing train output to: %s' % train_data_path_tsv)\n",
        "  print('Writing guide output to: %s' % guide_data_path_tsv)\n",
        "  with open(train_data_path_tsv, 'w') as outfile_train:\n",
        "    with open(guide_data_path_tsv, 'w') as outfile_guide:\n",
        "      csv_writer_train = csv.writer(outfile_train, delimiter='\\t')\n",
        "      csv_writer_guide = csv.writer(outfile_guide, delimiter='\\t')\n",
        "      for num_done_examples, example in enumerate(\n",
        "          tfds.load(tfds_name, split='train')):\n",
        "        if num_done_examples % LOGGING_STEPS == 0:\n",
        "          print('%d train examples done.' % num_done_examples)\n",
        "        if split_for_guidance_data > 0 and num_done_examples % (\n",
        "            1 / split_for_guidance_data) == 0:\n",
        "          csv_writer_guide.writerow([\n",
        "              _read_string_tensor(example[source_language]),\n",
        "              _read_string_tensor(example[target_language])\n",
        "          ])\n",
        "          guide_example_count += 1\n",
        "        else:\n",
        "          csv_writer_train.writerow([\n",
        "              _read_string_tensor(example[source_language]),\n",
        "              _read_string_tensor(example[target_language])\n",
        "          ])\n",
        "          train_example_count += 1\n",
        "\n",
        "  print('Num train examples: %d' % train_example_count)\n",
        "  print('Num guide examples: %d' % guide_example_count)\n",
        "  dev_data_path_tsv = os.path.join(output_dir, '{}_dev.tsv'.format(lang_pair))\n",
        "\n",
        "  with open(dev_data_path_tsv, 'w') as outfile_dev:\n",
        "    csv_writer_dev = csv.writer(outfile_dev, delimiter='\\t')\n",
        "    for num_done_examples, example in enumerate(\n",
        "        tfds.load(tfds_name, split='validation')):\n",
        "      csv_writer_dev.writerow([\n",
        "          _read_string_tensor(example[source_language]),\n",
        "          _read_string_tensor(example[target_language])\n",
        "      ])\n",
        "    print('Num validation examples: %d' % num_done_examples)\n",
        "\n",
        "  return train_data_path_tsv, guide_data_path_tsv, dev_data_path_tsv\n",
        "\n",
        "def generate_vocab(\n",
        "    output_dir,\n",
        "    source_language,\n",
        "    target_language,\n",
        "    tfds_name,\n",
        "):\n",
        "  \"\"\"Train a sentencepiece vocab on a portion of the data.\n",
        "\n",
        "  Args:\n",
        "    output_dir: The dir to which the data will be written. Dirs will be\n",
        "      recursively created if not already present.\n",
        "    source_language: Source language of the translation data.\n",
        "    target_language: Target language of the translation data.\n",
        "    tfds_name: The name of the desired dataset in tfds. (ie.\n",
        "      wmt19_translate/de-en). See\n",
        "      https://www.tensorflow.org/datasets/catalog/wmt19_translate for more\n",
        "        details.\n",
        "  \"\"\"\n",
        "\n",
        "  tf.io.gfile.makedirs(output_dir)\n",
        "  train_ds = tfds.load(tfds_name, split='train')\n",
        "\n",
        "  spm_model_path = os.path.join(\n",
        "      output_dir, '{}.32k.spm.model'.format(source_language + target_language))\n",
        "  print('spm_model_path %s' % spm_model_path)\n",
        "  _train_sentencepiece(\n",
        "      dataset=train_ds,\n",
        "      model_path=spm_model_path,\n",
        "      vocab_size=2**15,\n",
        "      data_keys=(source_language, target_language))\n",
        "  return spm_model_path\n",
        "\n",
        "\n",
        "def _dump_chars_to_textfile(dataset,\n",
        "                            maxchars=int(1e7),\n",
        "                            data_keys=('inputs', 'targets')):\n",
        "  \"\"\"Write part of a TFDS sentence dataset to lines in a text file.\n",
        "\n",
        "  Args:\n",
        "    dataset: tf.dataset containing string-data.\n",
        "    maxchars: int: approximate number of characters to save from dataset.\n",
        "    data_keys: Tuple[str]: what keys in dataset to dump from.\n",
        "\n",
        "  Returns:\n",
        "    name of temp file with dataset bytes, exact number of characters dumped.\n",
        "  \"\"\"\n",
        "  char_count = 0\n",
        "  ds_iter = dataset.as_numpy_iterator()\n",
        "  with tempfile.NamedTemporaryFile(\n",
        "      delete=False, prefix='/tmp/ds_chars') as outfp:\n",
        "    while char_count < maxchars:\n",
        "      example = next(ds_iter)\n",
        "      for k in data_keys:\n",
        "        line = example[k] + b'\\n'\n",
        "        char_count += len(line)\n",
        "        outfp.write(line)\n",
        "  return outfp.name, char_count\n",
        "\n",
        "\n",
        "def _train_sentencepiece(dataset,\n",
        "                         model_path,\n",
        "                         vocab_size=2**15,\n",
        "                         maxchars=int(1e7),\n",
        "                         model_type='unigram',\n",
        "                         character_coverage=1.0,\n",
        "                         data_keys=('inputs', 'targets')):\n",
        "  \"\"\"Train SentencePiece tokenizer from subset of tf dataset.\n",
        "\n",
        "  Args:\n",
        "    dataset: tf.dataset\n",
        "    model_path: str: path of model file to save vocab model to.\n",
        "    vocab_size: int: size of vocab tokens to train.\n",
        "    maxchars: int: number of characters to use for sentencepiece training.\n",
        "    model_type: str: type of sentencepiece vocab to train.\n",
        "    character_coverage: amount of characters covered by the model, good defaults\n",
        "      are 0.9995 for languages with rich character set like Japanese or Chinese\n",
        "      and 1.0 for other languages with small character set.\n",
        "    data_keys: Tuple[str]: keys of dataset to use for training.\n",
        "\n",
        "  Returns:\n",
        "    path to the trained sentencepiece vocabulary model.\n",
        "  \"\"\"\n",
        "  if model_path.startswith('gs://'):\n",
        "    abs_model_path = model_path\n",
        "  else:\n",
        "    abs_model_path = os.path.abspath(os.path.expanduser(model_path))\n",
        "  fname, _ = _dump_chars_to_textfile(\n",
        "      dataset, maxchars=maxchars, data_keys=data_keys)\n",
        "  with tempfile.NamedTemporaryFile(\n",
        "      delete=False, prefix='/tmp/sp_tmp') as model_fp:\n",
        "    pass  # we just want a prefix'd tmp-filename\n",
        "  argstr = ' '.join([\n",
        "      f'--input={fname}', f'--vocab_size={vocab_size}',\n",
        "      f'--character_coverage={character_coverage}',\n",
        "      f'--model_prefix={model_fp.name}', f'--model_type={model_type}'\n",
        "  ])\n",
        "  SentencePieceTrainer.Train(argstr)\n",
        "  # Use an intermediate filename that is renamed to the target name to address\n",
        "  # create and fill delays.\n",
        "  copy_rename_path = abs_model_path + '.rntmp'\n",
        "  tf.io.gfile.copy(model_fp.name + '.model', copy_rename_path, overwrite=True)\n",
        "  tf.io.gfile.rename(copy_rename_path, abs_model_path, overwrite=True)\n",
        "  print('copied %s to %s' % (model_fp.name + '.model', abs_model_path))\n",
        "  return abs_model_path\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "l0vDdVM587JE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Takes ~10 min with lt-en dataset.\n",
        "spm_model_path = generate_vocab(\n",
        "    output_dir=OUTPUT_DIR,\n",
        "    source_language=SOURCE_LANGUAGE,\n",
        "    target_language=TARGET_LANGUAGE,\n",
        "    tfds_name=TFDS_NAME,\n",
        ")"
      ],
      "metadata": {
        "id": "HDv-WJeL9HgT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Takes ~10 min with lt-en dataset (2321444 examples).\n",
        "train_data_path_tsv, guide_data_path_tsv, dev_data_path_tsv = write_data_to_tsv(\n",
        "    output_dir=OUTPUT_DIR,\n",
        "    source_language=SOURCE_LANGUAGE,\n",
        "    target_language=TARGET_LANGUAGE,\n",
        "    tfds_name=TFDS_NAME,\n",
        "    split_for_guidance_data=0.01,\n",
        ")"
      ],
      "metadata": {
        "id": "gJVDZpYdABax"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print('Done loading raw training data.\\n')\n",
        "\n",
        "print('vocab file: %s' % spm_model_path)\n",
        "print('train_data_path_tsv: %s' % train_data_path_tsv)\n",
        "print('guide_data_path_tsv: %s' % guide_data_path_tsv)\n",
        "print('dev_data_path_tsv: %s' % dev_data_path_tsv)"
      ],
      "metadata": {
        "id": "EVgbORmRE5_L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Produce training examples\n",
        "\n",
        "Convert raw examples to packed and spm-encoded tf.Examples. Takes about 5 min for lt-en dataset."
      ],
      "metadata": {
        "id": "dbHoE4sNB2h7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "\n",
        "INPUT_PATH=os.path.dirname(train_data_path_tsv)  # Input dir with TSV files.\n",
        "SPM_PATH=spm_model_path  # Path to the SPM model.\n",
        "\n",
        "\n",
        "TSV_SOURCE_COLUMN=0  # Source sentence TSV column.\n",
        "TSV_TARGET_COLUMN=1  # Target sentence TSV column. If negative, use all available columns except source_column.\n",
        "PACKED_LENGTH=256  # Length of packed examples. Set to zero to disable packing.\n",
        "PAD_LENGTH=256  # If positive, pad all features to this length.\n",
        "MIN_EDIT_DISTANCE=0.3  # Minimum edit distance. Examples for which source and target are more similar will be omitted.\n"
      ],
      "metadata": {
        "id": "Mu0116QhDH6k"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title\n",
        "\"\"\"A FlumePython program to generate tf.Examples with a sentencepiece model.\"\"\"\n",
        "\n",
        "_PAD_ID = -1\n",
        "_DEFAULT_EOS_ID = 2  # Only used if spm_path is None.\n",
        "\n",
        "class SelectTSVColumns(beam.DoFn):\n",
        "  \"\"\"Selects two columns in TSV lines.\"\"\"\n",
        "\n",
        "  def __init__(self, source_column=0, target_column=1):\n",
        "    self._source_column = source_column\n",
        "    self._target_column = target_column\n",
        "\n",
        "  def process(self, tsv_line):\n",
        "    columns = tsv_line.split('\\t')\n",
        "    try:\n",
        "      source_sentence = columns[self._source_column]\n",
        "    except IndexError:\n",
        "      metrics.Metrics.counter('err_malformed_tsv_source', 'count').inc()\n",
        "      return\n",
        "    if self._target_column >= 0:\n",
        "      try:\n",
        "        yield source_sentence, columns[self._target_column]\n",
        "      except IndexError:\n",
        "        metrics.Metrics.counter('err_malformed_tsv_target', 'count').inc()\n",
        "    else:\n",
        "      for column_index, target_sentence in enumerate(columns):\n",
        "        if column_index != self._source_column:\n",
        "          yield source_sentence, target_sentence\n",
        "\n",
        "\n",
        "class PrepareTfExamples(beam.DoFn):\n",
        "  \"\"\"Prepare (packed) TFExamples from a list of source/target sentence pairs.\"\"\"\n",
        "\n",
        "  def __init__(self, spm_path, packed_length=256, pad_length=256):\n",
        "    self._spm_path = spm_path\n",
        "    self._spm = None\n",
        "    self._packed_length = packed_length\n",
        "    self._packed_examples = packed_length > 0\n",
        "    self._pad_length = pad_length\n",
        "\n",
        "  def start_bundle(self):\n",
        "    if self._spm_path:\n",
        "      with tf.io.gfile.GFile(self._spm_path, 'rb') as f:\n",
        "        spm_model = f.read()\n",
        "      self._spm = spm.SentencePieceProcessor()\n",
        "      self._spm.LoadFromSerializedProto(spm_model)\n",
        "\n",
        "  def _make_spm_example_dict(self, source_text, target_text):\n",
        "    return {\n",
        "        'inputs': self._encode_with_spm(source_text),\n",
        "        'targets': self._encode_with_spm(target_text)\n",
        "    }\n",
        "\n",
        "  def _encode_with_spm(self, text):\n",
        "    if self._spm is not None:\n",
        "      return self._spm.EncodeAsIds(text) + [self._spm.eos_id()]\n",
        "    return [int(t) for t in text.strip().split()] + [_DEFAULT_EOS_ID]\n",
        "\n",
        "  def _pad_example_dict(self, example_dict):\n",
        "    if self._pad_length <= 0:\n",
        "      return example_dict\n",
        "\n",
        "    padded_example_dict = {}\n",
        "    for key, sequence in example_dict.items():\n",
        "      num_pads = self._pad_length - len(sequence)\n",
        "      if num_pads < 0:\n",
        "        raise ValueError('Feature %r too long' % key)\n",
        "\n",
        "      padded_example_dict[key] = sequence + [_PAD_ID] * num_pads\n",
        "    return padded_example_dict\n",
        "\n",
        "  def process(self, source_target_list):\n",
        "    example_dicts = [\n",
        "        self._make_spm_example_dict(source_text, target_text)\n",
        "        for source_text, target_text in source_target_list\n",
        "    ]\n",
        "    if self._packed_examples:\n",
        "      example_dicts = pack_examples(\n",
        "          example_dicts, has_inputs=True, packed_length=self._packed_length)\n",
        "    for example_dict in example_dicts:\n",
        "      try:\n",
        "        padded_example_dict = self._pad_example_dict(example_dict)\n",
        "      except ValueError:\n",
        "        metrics.Metrics.counter('err_too_long', 'count').inc()\n",
        "      else:\n",
        "        yield to_example(padded_example_dict).SerializeToString()\n",
        "\n",
        "\n",
        "class ValidateSentencePair(beam.DoFn):\n",
        "  \"\"\"String transformation applied to a sequence pair to validate the quality of a sample.\"\"\"\n",
        "\n",
        "  def __init__(self, min_distance):\n",
        "    self.min_distance = min_distance\n",
        "\n",
        "  def validate_similarity(self, source, target):\n",
        "    \"\"\"Discard sentence pair if their distance is less or equal than the threshold.\"\"\"\n",
        "    distance = difflib.SequenceMatcher(None, source, target).ratio()\n",
        "    return distance > self.min_distance\n",
        "\n",
        "  def process(self, sequence_pair):\n",
        "    stripped_sequence_pair = sequence_pair.strip()\n",
        "\n",
        "    try:\n",
        "      source_sequence, target_sequence = stripped_sequence_pair.split('\\t')\n",
        "    except ValueError:\n",
        "      metrics.Metrics.counter('erroneous_tab_split', 'count').inc()\n",
        "      return\n",
        "\n",
        "    source_tokens = source_sequence.split()\n",
        "    target_tokens = target_sequence.split()\n",
        "\n",
        "    if not self.validate_similarity(source_tokens, target_tokens):\n",
        "      logging.info('Discarded due to high similarity: %r, %r', source_sequence,\n",
        "                   target_sequence)\n",
        "      metrics.Metrics.counter('high_similarity', 'count').inc()\n",
        "      return\n",
        "\n",
        "    metrics.Metrics.counter('samples_processed', 'count').inc()\n",
        "    yield stripped_sequence_pair\n",
        "\n",
        "\n",
        "def pipeline(root):\n",
        "  \"\"\"Method to pass into flume runner.\"\"\"\n",
        "  for i, tsv_in in enumerate(tf.io.gfile.glob(INPUT_PATH+'/*.tsv')):\n",
        "    print('Processing tsv input: %s' % tsv_in)\n",
        "    tfr_out = tsv_in.replace('.tsv', '.tfr')\n",
        "    _ = (\n",
        "        root\n",
        "        | 'Read RecordIO TSV__%s' % i >> beam.io.ReadFromText(tsv_in)\n",
        "        | 'Validate sentence pair__%s' % i >> beam.ParDo(\n",
        "            ValidateSentencePair(MIN_EDIT_DISTANCE))\n",
        "        | 'Select TSV columns__%s' % i >> beam.ParDo(\n",
        "            SelectTSVColumns(\n",
        "                source_column=TSV_SOURCE_COLUMN,\n",
        "                target_column=TSV_TARGET_COLUMN))\n",
        "        | 'Reshuffle__%s' % i >> beam.Reshuffle()\n",
        "        | 'Batch elements__%s' % i >> beam.BatchElements(\n",
        "            min_batch_size=1024, max_batch_size=1024)\n",
        "        | 'Make tf.Examples__%s' % i >> beam.ParDo(\n",
        "            PrepareTfExamples(\n",
        "                spm_path=SPM_PATH,\n",
        "                packed_length=PACKED_LENGTH,\n",
        "                pad_length=PAD_LENGTH))\n",
        "        | 'Write to tf.Record__%s' % i >> tfrecordio.WriteToTFRecord(tfr_out))\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "xC2s7DEGB1hW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "\"\"\"Runs the Beam pipeline, can take a while. \"\"\"\n",
        "# ~3 min for lt-en dataset\n",
        "\n",
        "options = pipeline_options.PipelineOptions()\n",
        "p = beam.Pipeline(options=options)\n",
        "pipeline(p)\n",
        "p.run().wait_until_finish()\n"
      ],
      "metadata": {
        "id": "Fs3xiCUEEqBs"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Training\n",
        "\n",
        "Code, setup, and demonstration runs for training."
      ],
      "metadata": {
        "id": "b6iUGKmdOl-k"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load Code\n",
        "\n",
        "Runs imports and loads training code (copied from train.py rather than loaded from the github to allow for easy editing). If you edit the code be sure to re-run this cell block before training to pick up the changes."
      ],
      "metadata": {
        "id": "Hd0GYpo5jsD4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Imports\n",
        "\"\"\"Training script for gradient-based hyper-parameter tuning in JAX.\"\"\"\n",
        "\n",
        "\n",
        "import functools\n",
        "import json\n",
        "import os\n",
        "import textwrap\n",
        "import time\n",
        "\n",
        "from absl import app\n",
        "from absl import flags\n",
        "from absl import logging\n",
        "from flax import jax_utils\n",
        "from flax import serialization\n",
        "from flax.deprecated import nn\n",
        "from flax.metrics import tensorboard\n",
        "from flax.training import checkpoints\n",
        "from flax.training import common_utils\n",
        "import jax\n",
        "from jax import random\n",
        "import jax.nn\n",
        "import jax.numpy as jnp\n",
        "import sentencepiece as spm\n",
        "import tensorflow.compat.v2 as tf\n",
        "from gradient_based_tuning import data\n",
        "from gradient_based_tuning import guided_parameters\n",
        "from gradient_based_tuning import models\n",
        "from gradient_based_tuning import utils\n",
        "\n",
        "%load_ext tensorboard"
      ],
      "metadata": {
        "id": "0-Me1BgGJVaA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Define Training FLAGS\n",
        "\n",
        "\n",
        "if not 'model_dir' in flags.FLAGS:  # Do not try to define FLAGS twice.\n",
        "  FLAGS = flags.FLAGS\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'model_dir', default=None, help='Directory to store model data.')\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'vocab_path',\n",
        "      default=None,\n",
        "      help='Path to load or store word piece vocab file.',\n",
        "      required=True)\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'jax_random_seed', default=0, help='Integer for PRNG random seed.')\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'jax_backend_target',\n",
        "      default=None,\n",
        "      help=('TPU grpc target for use with cloud TPUs.'\n",
        "            ' e.g. grpc://192.168.0.2:8470'))\n",
        "\n",
        "  # Training flags.\n",
        "  flags.DEFINE_string(\n",
        "      'training_dataset',\n",
        "      default='lang8',\n",
        "      help='Which dataset to use for training.')\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'training_dataset_path',\n",
        "      default=None,\n",
        "      help='Path to prepacked TFRecords for training.',\n",
        "      required=True)\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'guidance_dataset_path',\n",
        "      default='',\n",
        "      help='Path to prepacked TFRecords for guidance.')\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'eval_dataset_path',\n",
        "      default='',\n",
        "      help='Path to prepacked TFRecords for evaluation.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'batch_size', default=256, help='Per host batch size for training.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'num_train_steps',\n",
        "      default=160000,\n",
        "      help='Number of train steps. Total train '\n",
        "      'steps = min(num_train_steps, max_train_epochs * <steps per epoch>).')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'max_train_epochs',\n",
        "      default=100,\n",
        "      help='Maximum bound of how many epochs to train, will cut training short if'\n",
        "      ' this number is reached before num_train_steps.')\n",
        "\n",
        "  flags.DEFINE_float('learning_rate', default=0.4, help='Base learning rate.')\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'learning_rate_schedule',\n",
        "      default='constant * linear_warmup * rsqrt_decay',\n",
        "      help='Schedule for learning rate of model.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'warmup_steps_ratio',\n",
        "      default=0.1,\n",
        "      help='Proportion of FLAGS.num_train_steps to spend on warmup. '\n",
        "      'Overridden by FLAGS.warmup_steps.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'warmup_steps',\n",
        "      default=0,\n",
        "      help='How many steps to run linear learning rate warmup. If non-zero, '\n",
        "      'overrides warmup_steps_ratio.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'guide_with_train_loss',\n",
        "      default=False,\n",
        "      help='If True, use train batch for guide batch')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'reset_guidance_per_steps',\n",
        "      default=0,\n",
        "      help='if nonzero, reset all guided vars every X steps')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'max_target_length',\n",
        "      default=256,\n",
        "      help='Maximum length cutoff for training/guide/eval examples.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'use_model_lr_for_guided',\n",
        "      default=False,\n",
        "      help='If True, use the default model LR schedule for guided vars.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'save_checkpoints', default=True, help='Whether to save model checkpoints.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'keep_checkpoints_count',\n",
        "      default=10,\n",
        "      help='How many model checkpoints to keep around at a time.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'save_checkpoint_at_init',\n",
        "      default=True,\n",
        "      help='Whether to save model checkpoint 0 at initialization.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'clip_nan_grads',\n",
        "      default=True,\n",
        "      help='If True, clip NaNs in the hyper-parameter gradient.')\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'init_checkpoint',\n",
        "      default=None,\n",
        "      help='Pretrained checkpoint to initialize training.')\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'guided_hparam_types',\n",
        "      default='learning_rate',\n",
        "      help='Only used when \"param\" in model_optimizer_type, and train_with_guided_parameters is true.'\n",
        "  )\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'model_ckpt_min_freq',\n",
        "      default=10000,\n",
        "      help='Checkpoint model every X steps during training. If 0, ignore.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'train_metrics_per_eval_metric',\n",
        "      default=10,\n",
        "      help='Emit train metrics this many times per emitting eval metrics.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'save_ckpt_per_epoch',\n",
        "      default=False,\n",
        "      help='If True, save a ckpt per epoch. Supersedes model_ckpt_min_freq.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'label_smoothing', default=0.0, help='Cross entropy loss label smoothing.')\n",
        "\n",
        "  flags.DEFINE_float('dropout_rate', default=0.1, help='Dropout rate.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'attention_dropout_rate', default=0.1, help='Attention dropout rate.')\n",
        "\n",
        "  # Eval flags.\n",
        "  flags.DEFINE_boolean(\n",
        "      'do_eval', default=True, help='Whether or not to perform evaluation.')\n",
        "\n",
        "  flags.DEFINE_list(\n",
        "      'eval_dataset_list',\n",
        "      default=[\n",
        "          'bea_dev',\n",
        "          'conll14',\n",
        "          'fce_test',\n",
        "      ],\n",
        "      help=textwrap.dedent(\"\"\"Which datasets to evaluate on. Options:\n",
        "          bea_dev,\n",
        "          bea_dev_abc_half1,\n",
        "          bea_dev_abc_half2,\n",
        "          bea_dev_n_half1,\n",
        "          bea_dev_n_half2,\n",
        "          bea_train,\n",
        "          conll13,\n",
        "          conll14,\n",
        "          fce_test,\n",
        "          fce_train,\n",
        "          uxr.\"\"\"),\n",
        "  )\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'eval_batch_size', default=16, help='Per host batch size for eval.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'min_eval_freq',\n",
        "      default=500,\n",
        "      help='Minimum bound of evaluation frequency during model training.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'max_metric_freq',\n",
        "      default=250,\n",
        "      help='Minimum bound of how many evaluations to make during model training.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'min_train_metric_freq',\n",
        "      default=250,\n",
        "      help='Minimum bound of how many evaluations to make during model training.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'num_eval_steps',\n",
        "      default=1024,\n",
        "      help='Number of steps to take during evaluation.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'total_evals',\n",
        "      default=200,\n",
        "      help='Eval at least this many times over the course of training. '\n",
        "      'Can be overridden by FLAGS.min_eval_freq.')\n",
        "\n",
        "  # Guided learning flags.\n",
        "  flags.DEFINE_bool(\n",
        "      'train_with_guided_parameters',\n",
        "      default=True,\n",
        "      help='If True, train with guided parameters. Else, vanilla training.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'guided_weight_decay',\n",
        "      default=0,\n",
        "      help='weight_decay decay for guided optimizers')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'guided_ckpt_min_freq',\n",
        "      default=0,\n",
        "      help='Checkpoint guided optimizers every X steps during training. Upper-'\n",
        "      'bounded by model_ckpt_min_freq. If 0, fall back to model_ckpt_min_freq.'\n",
        "      'The only reason to have this at greater frequency that model_ckpt_min_freq'\n",
        "      ' is for guided learning visualization purposes.')\n",
        "\n",
        "  flags.DEFINE_enum(\n",
        "      'guide_batch_update_freq', 'NEVER',\n",
        "      ['NEVER', 'PER_EPOCH', 'PER_TRAINING_STEP', 'PER_GUIDANCE_STEP'],\n",
        "      'How often to update the guidance batch.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'guide_batch_size', default=256, help='Per host batch size for guidance.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'learning_rate_guided_vars',\n",
        "      default=3e-5,\n",
        "      help='Base learning rate for guided parameters.')\n",
        "\n",
        "  flags.DEFINE_string(\n",
        "      'learning_rate_schedule_guided_vars',\n",
        "      default='constant',\n",
        "      help='Schedule for learning rate of guided parameters.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'warmup_steps_ratio_guided_vars',\n",
        "      default=0.1,\n",
        "      help='Linear learning rate warmup for guided parameters.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'grad_clip_limit',\n",
        "      default=1.0,\n",
        "      help='0.0 means no clip limit. 10 is reasonable, if NaN, try lower.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'use_grad_clip',\n",
        "      default=False,\n",
        "      help='If True, apply FLAGS.grad_clip_limit to guided parameter gradients.')\n",
        "\n",
        "  # Decode/predict flags.\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'jax_beam_size', default=4, help='Beam size for inference.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'max_predict_length',\n",
        "      default=64,\n",
        "      help='Maximum length cutoff for predicted tokens. Predicted examples are '\n",
        "      'not packed so this is max length for an individual example.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'do_predict',\n",
        "      default=True,\n",
        "      help='Whether to run beam predictions on the eval set.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'do_score',\n",
        "      default=False,\n",
        "      help='Whether to run beam predictions on the eval set.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'num_decoding_iterations', default=2, help='Number of decoding iterations.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'identity_penalty', default=1.0, help='Identity penalty factor.')\n",
        "\n",
        "  flags.DEFINE_list('testsets', default=[], help='Which testsets to run.')\n",
        "\n",
        "  flags.DEFINE_list(\n",
        "      'identity_penalty_list',\n",
        "      default=[0.8, 1, 2.0, 3.0],\n",
        "      help='which identity_penalty values to grid search, override with csv')\n",
        "\n",
        "  # Model flags\n",
        "\n",
        "  flags.DEFINE_enum('model_optimizer_type', 'lamb', [\n",
        "      'lamb',\n",
        "      'gd',\n",
        "      'adam',\n",
        "      'sgd',\n",
        "      'mom',\n",
        "      'momentum',\n",
        "      'adagrad',\n",
        "      'adafactor',\n",
        "  ], 'Which optimizer to use for model optimization.')\n",
        "\n",
        "  flags.DEFINE_enum(\n",
        "      'guided_params_optimizer_type', 'adam', [\n",
        "          'lamb',\n",
        "          'gd',\n",
        "          'adam',\n",
        "          'sgd',\n",
        "          'mom',\n",
        "          'momentum',\n",
        "          'adagrad',\n",
        "          'adafactor',\n",
        "      ], 'Which optimizer to use for guided parameter meta-optimization.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'share_embeddings',\n",
        "      default=True,\n",
        "      help='Inputs and targets share embedding.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'logits_via_embedding',\n",
        "      default=True,\n",
        "      help='Final logit transform uses embedding matrix transpose.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'num_layers', default=6, help='Number of transformer layers.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'qkv_dim', default=1024, help='Size of query/key/value for attention.')\n",
        "\n",
        "  flags.DEFINE_integer('emb_dim', default=1536, help='Size of embeddings.')\n",
        "\n",
        "  flags.DEFINE_integer('mlp_dim', default=4096, help='Size of the MLP.')\n",
        "\n",
        "  flags.DEFINE_integer('num_heads', default=4, help='Number of attention heads.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'use_bfloat16',\n",
        "      default=True,\n",
        "      help=('Use bfloat16 mixed precision training instead of float32.'))\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'take_current_run_step_from_init',\n",
        "      default=False,\n",
        "      help=('If True, when init from ckpt_X, current run step set to X. Else 0.'))\n",
        "\n",
        "  # Model optimizer flags - only those for FLAGS.model_optimizer_type will be used\n",
        "\n",
        "  flags.DEFINE_float('adagrad_eps', default=1e-8, help='Adagrad epsilon.')\n",
        "\n",
        "  flags.DEFINE_float('momentum_beta1', default=0.9, help='Momentum beta1.')\n",
        "  flags.DEFINE_float(\n",
        "      'momentum_weight_decay', default=0.0, help='Momentum weight_decay.')\n",
        "  flags.DEFINE_bool(\n",
        "      'momentum_nesterov', default=True, help='Momentum use nesterov.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'adafactor_factored',\n",
        "      default=True,\n",
        "      help='Adafactor adafactor_multiply_by_parameter_scale')\n",
        "  flags.DEFINE_bool(\n",
        "      'adafactor_multiply_by_parameter_scale',\n",
        "      default=True,\n",
        "      help='Adafactor adafactor_multiply_by_parameter_scale')\n",
        "  flags.DEFINE_float(\n",
        "      'adafactor_beta1', default=None, help='Adafactor optimizer beta1.')\n",
        "  flags.DEFINE_float(\n",
        "      'adafactor_epsilon1', default=1e-30, help='Adafactor optimizer eps1.')\n",
        "  flags.DEFINE_float(\n",
        "      'adafactor_epsilon2', default=1e-3, help='Adafactor optimizer eps2.')\n",
        "  flags.DEFINE_float(\n",
        "      'adafactor_decay_rate', default=0.8, help='Adafactor optimizer decay rate.')\n",
        "  flags.DEFINE_integer(\n",
        "      'adafactor_step_offset',\n",
        "      default=0,\n",
        "      help='Adafactor adafactor_multiply_by_parameter_scale')\n",
        "  flags.DEFINE_float(\n",
        "      'adafactor_clipping_threshold',\n",
        "      default=1.0,\n",
        "      help='Adafactor adafactor_multiply_by_parameter_scale')\n",
        "  flags.DEFINE_float(\n",
        "      'adafactor_weight_decay_rate',\n",
        "      default=None,\n",
        "      help='Adafactor adafactor_multiply_by_parameter_scale')\n",
        "  flags.DEFINE_integer(\n",
        "      'adafactor_min_dim_size_to_factor',\n",
        "      default=128,\n",
        "      help='Adafactor adafactor_multiply_by_parameter_scale')\n",
        "\n",
        "  flags.DEFINE_float('adam_beta1', default=0.9, help='Adam optimizer beta1.')\n",
        "  flags.DEFINE_float('adam_beta2', default=0.999, help='Adam optimizer beta2.')\n",
        "  flags.DEFINE_float(\n",
        "      'adam_beta3', default=0.1, help='smAdam decay optimizer beta3.')\n",
        "  flags.DEFINE_float('adam_eps', default=1e-6, help='Adam optimizer epsilon.')\n",
        "  flags.DEFINE_float(\n",
        "      'adam_weight_decay', default=0.0, help='Adam optimizer weight_decay.')\n",
        "\n",
        "  flags.DEFINE_float('lamb_beta1', default=0.9, help='LAMB optimizer beta1.')\n",
        "  flags.DEFINE_float('lamb_beta2', default=0.999, help='LAMB optimizer beta2.')\n",
        "  flags.DEFINE_float('lamb_eps', default=1e-6, help='LAMB optimizer epsilon.')\n",
        "  flags.DEFINE_float(\n",
        "      'lamb_weight_decay', default=0.0, help='LAMB optimizer weight_decay.')\n",
        "\n",
        "  # pylint: enable=line-too-long\n",
        "  flags.DEFINE_float(\n",
        "      'weight_decay',\n",
        "      default=0,\n",
        "      help='If the model_optimizer_type supports weight decay, it '\n",
        "      'will be initialized by this value. May be overridden by '\n",
        "      'guided hparam.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'ghp_lrs_beta1',\n",
        "      default=1.0,\n",
        "      help='Scale the meta-optimizer guided learning rate for beta1.')\n",
        "  flags.DEFINE_float(\n",
        "      'ghp_lrs_eps',\n",
        "      default=1.0,\n",
        "      help='Scale the meta-optimizer guided learning rate for epsilon.')\n",
        "  flags.DEFINE_float(\n",
        "      'ghp_lrs_lr',\n",
        "      default=1.0,\n",
        "      help='Scale the meta-optimizer guided learning rate for learning_rate.')\n",
        "  flags.DEFINE_float(\n",
        "      'ghp_lrs_wd',\n",
        "      default=1.0,\n",
        "      help='Scale the meta-optimizer guided learning rate for weight_decay.')\n",
        "  flags.DEFINE_float(\n",
        "      'ghp_lrs_ls',\n",
        "      default=1.0,\n",
        "      help='Scale the meta-optimizer guided learning rate for label_smoothing.')\n",
        "\n",
        "  flags.DEFINE_float(\n",
        "      'guided_update_beta1',\n",
        "      default=None,\n",
        "      help='Use to override meta-optimizer default beta1, if applicable.')\n",
        "  flags.DEFINE_float(\n",
        "      'guided_update_beta2',\n",
        "      default=None,\n",
        "      help='Use to override meta-optimizer default beta2, if applicable.')\n",
        "  flags.DEFINE_float(\n",
        "      'guided_update_eps',\n",
        "      default=None,\n",
        "      help='Use to override meta-optimizer default epsilon, if applicable.')\n",
        "  flags.DEFINE_float(\n",
        "      'guided_update_wd',\n",
        "      default=None,\n",
        "      help='Use to override meta-optimizer default weight_decay, if applicable.')\n",
        "\n",
        "  flags.DEFINE_bool(\n",
        "      'print_for_colab',\n",
        "      default=False,\n",
        "      help='If true, print to console as well as output logging.')\n",
        "\n",
        "  flags.DEFINE_integer(\n",
        "      'max_eval_steps',\n",
        "      default=256,\n",
        "      help='Max number of steps to take during evaluation.')\n",
        "\n",
        "_TRAIN_KEYS_BASE = [\n",
        "    'inputs',\n",
        "    'targets',\n",
        "    'inputs_position',\n",
        "    'targets_position',\n",
        "    'inputs_segmentation',\n",
        "    'targets_segmentation',\n",
        "]\n",
        "_TRAINING_DONE_FILENAME = 'training_done.txt'\n"
      ],
      "metadata": {
        "id": "ovo1LPyiJUol"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "jax.devices()"
      ],
      "metadata": {
        "id": "7pocXjnoelvV"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Guidance Step Code\n",
        "def update_model_and_guidance_step(\n",
        "    optimizer_dict,\n",
        "    train_batch,\n",
        "    guided_vars_dict,\n",
        "    learning_rate_fn,\n",
        "    learning_rate_fn_guided_vars,\n",
        "    guide_batch,\n",
        "    grad_clip_limit=0.0,\n",
        "    use_bfloat16=False,\n",
        "    dropout_rng=None,\n",
        "    output_metrics=None,\n",
        "):\n",
        "  \"\"\"Updates all guided optimizers and the model, with a single forward pass.\n",
        "\n",
        "  If multiple guided optimizers are in use, they are updated simultaneously.\n",
        "\n",
        "  Args:\n",
        "    optimizer_dict: a dictionary holding all optimizers\n",
        "    train_batch: the batch of training data to be applied\n",
        "    guided_vars_dict: specifies the guided parameter configuration\n",
        "    learning_rate_fn: model learning rate function\n",
        "    learning_rate_fn_guided_vars: guided vars learning rate function\n",
        "    guide_batch: the batch of guidance data to be applied\n",
        "    grad_clip_limit: if nonzero, clip all gradients to abs magnitude <= limit\n",
        "    use_bfloat16: bool, If True, use lower-precision bfloat16\n",
        "    dropout_rng: the JAX rng for dropout\n",
        "    output_metrics: dict of metrics to write out to tensorboard\n",
        "\n",
        "  Returns:\n",
        "    optimizer_dict: updated with steps taken for guided parameter optimizers\n",
        "    metrics: metrics dict keeping track of tensorboard-reported summary values\n",
        "    new_dropout_rng: updated RNG\n",
        "  \"\"\"\n",
        "  metrics_joint = {}  # for summarizing\n",
        "  model_optimizer = optimizer_dict['model']  # for readability\n",
        "  dropout_rng, new_dropout_rng = random.split(dropout_rng)\n",
        "  current_step = get_step_from_opt(model_optimizer)\n",
        "\n",
        "  # Extract fields of interest from the training batch.\n",
        "  (inputs, targets, inputs_positions, targets_positions, inputs_segmentation,\n",
        "   targets_segmentation) = [train_batch.get(k, None) for k in _TRAIN_KEYS_BASE]\n",
        "\n",
        "  # Extract fields from the guidance batch.\n",
        "  (guide_inputs, guide_targets, guide_inputs_positions, guide_targets_positions,\n",
        "   guide_inputs_segmentation, guide_targets_segmentation) = [\n",
        "       guide_batch.get(k, None) for k in _TRAIN_KEYS_BASE\n",
        "   ]\n",
        "\n",
        "  # Organize raw vars and activation functions for easy access.\n",
        "  raw_vars_dict, act_fn_dict = guided_parameters.get_raw_vars_and_act_fns(\n",
        "      optimizer_dict, guided_vars_dict)\n",
        "\n",
        "  # Learning rate for the guided variables.\n",
        "  lr_guided_vars = learning_rate_fn_guided_vars(current_step)\n",
        "  metrics_joint['meta_hparams/learning_rate_for_guided_vars'] = lr_guided_vars\n",
        "  for x in guided_vars_dict:\n",
        "    metrics_joint[\n",
        "        'meta_hparams/learning_rate_for_guided_vars__%s' %\n",
        "        x.split('-')[-1].split('__')\n",
        "        [-1]] = lr_guided_vars * guided_vars_dict[x]['learning_rate_scalar']\n",
        "\n",
        "  # Loss fn for model step lookahead, used within loss_fn_guidance_step.\n",
        "  # Applies example weighting, smoothing, etc, to the lookahead model.\n",
        "  def loss_fn_model_lookahead(\n",
        "      model,\n",
        "      raw_label_smoothing,\n",
        "  ):\n",
        "    \"\"\"Loss fn for model step lookahead, used within loss_fn_guidance_step.\"\"\"\n",
        "\n",
        "    # Default to FLAGS value if not set within guided_vars_dict.\n",
        "    dropout_rate = FLAGS.dropout_rate\n",
        "    attention_dropout_rate = FLAGS.attention_dropout_rate\n",
        "    label_smoothing = FLAGS.label_smoothing if FLAGS.label_smoothing else 0.0\n",
        "\n",
        "    # It is necessary to do None comparison rather than ternary operator here\n",
        "    # because JAX will have tracer errors with the ternary operator.\n",
        "    if raw_label_smoothing is not None:\n",
        "      label_smoothing = act_fn_dict['hp-label_smoothing'](raw_label_smoothing)\n",
        "\n",
        "    with nn.stochastic(dropout_rng):\n",
        "      train_logits = model(\n",
        "          inputs=inputs,\n",
        "          targets=targets,\n",
        "          use_bfloat16=use_bfloat16,\n",
        "          inputs_positions=inputs_positions,\n",
        "          targets_positions=targets_positions,\n",
        "          inputs_segmentation=inputs_segmentation,\n",
        "          targets_segmentation=targets_segmentation,\n",
        "          train=True,\n",
        "          cache=None,\n",
        "          dropout_rate=dropout_rate,\n",
        "          attention_dropout_rate=attention_dropout_rate)\n",
        "\n",
        "    # Initialize weights using all available guided parameters.\n",
        "    weights = jnp.where(targets > 0, 1, 0).astype(jnp.float32)\n",
        "\n",
        "    # Now that all dp-vars have been applied to weights, calculate loss.\n",
        "    loss, weight_sum = utils.compute_weighted_cross_entropy(\n",
        "        train_logits,\n",
        "        targets,\n",
        "        weights,\n",
        "        label_smoothing,\n",
        "    )\n",
        "    full_train_loss = loss / jnp.where(weight_sum != 0, weight_sum, 1)\n",
        "\n",
        "    return full_train_loss, (train_logits,)\n",
        "\n",
        "  # Loss fn for guided variables, contains a call to loss_fn_model_lookahead.\n",
        "  def loss_fn_guidance_step(\n",
        "      model_opt,\n",
        "      map_guided_names_to_raw_indices,\n",
        "      *per_param_raw_values,\n",
        "  ):\n",
        "    \"\"\"Loss function used for guided parameter variables.\n",
        "\n",
        "    Summary:\n",
        "      this returns validation loss of the next-step model as a function of the\n",
        "      current model and the current weights.\n",
        "\n",
        "      The partial derivative of this function with respect to the raw_vars gives\n",
        "      the gradient direction for the updates to the raw vars, as determined by\n",
        "      their marginal contribution to the loss on the validation set. It is\n",
        "      necessary to pass the raw vars (as opposed to the optimizer that wraps\n",
        "      them) in order to take the gradient of this function with respect to the\n",
        "      vars themselves.\n",
        "\n",
        "    Steps:\n",
        "      calculate current model next train step gradient using current weights\n",
        "      apply gradient to get 'next-step' model\n",
        "      calculate validation loss for 'next-step' model\n",
        "      (model_opt_lookahead.target)\n",
        "      return mean_val_loss\n",
        "\n",
        "    Args:\n",
        "      model_opt: optimizer that wraps Transformer model. the model is accessed\n",
        "        via model_opt.target. It is necessary to pass the optimizer in (rather\n",
        "        than just the model itself) so that the model may be updated via the\n",
        "        optimizer's model_opt.apply_gradient(...) fn.\n",
        "      map_guided_names_to_raw_indices: dict mapping guided names to raw value\n",
        "        indices within per_param_raw_values\n",
        "      *per_param_raw_values: raw values of guided params, with indices specified\n",
        "        by map_guided_names_to_raw_indices\n",
        "\n",
        "    Returns:\n",
        "      mean_val_loss\n",
        "      guide_logits\n",
        "    \"\"\"\n",
        "    global_guided_vars_count = 0\n",
        "\n",
        "    raw_label_smoothing = None\n",
        "    if 'hp-label_smoothing' in map_guided_names_to_raw_indices:\n",
        "      raw_label_smoothing = per_param_raw_values[\n",
        "          map_guided_names_to_raw_indices['hp-label_smoothing']]\n",
        "      global_guided_vars_count += 1\n",
        "\n",
        "    # All raw guided vars need to be activated.\n",
        "    model_learning_rate = learning_rate_fn(current_step)\n",
        "    if 'hp-learning_rate' in map_guided_names_to_raw_indices:\n",
        "      activated_lr = act_fn_dict['hp-learning_rate'](\n",
        "          per_param_raw_values[\n",
        "              map_guided_names_to_raw_indices['hp-learning_rate']])\n",
        "      model_learning_rate *= activated_lr\n",
        "      global_guided_vars_count += 1\n",
        "\n",
        "    # calculate current model next train step gradient using current weights\n",
        "    (full_train_loss,\n",
        "     (train_logits,)), model_train_step_grad = jax.value_and_grad(\n",
        "         loss_fn_model_lookahead, argnums=0, has_aux=True)(\n",
        "             model_opt.target,\n",
        "             raw_label_smoothing=raw_label_smoothing,\n",
        "         )\n",
        "    model_train_step_grad = jax.lax.pmean(model_train_step_grad, 'batch')\n",
        "\n",
        "    opt_kwargs_dict = _get_default_opt_kwargs(\n",
        "        model_optimizer_type=FLAGS.model_optimizer_type,\n",
        "        lr=model_learning_rate,\n",
        "    )\n",
        "    # Override defaults with learned values.\n",
        "    for hparam in opt_kwargs_dict:\n",
        "      if 'hp-%s' % hparam in map_guided_names_to_raw_indices:\n",
        "        if hparam == 'learning_rate':  # Skip LR as it is already set above.\n",
        "          continue\n",
        "        opt_kwargs_dict[hparam] = act_fn_dict['hp-%s' % hparam](\n",
        "            per_param_raw_values[map_guided_names_to_raw_indices['hp-%s' %\n",
        "                                                                 hparam]])\n",
        "    apply_grad_ret_list = model_opt.apply_gradient(\n",
        "        model_train_step_grad,\n",
        "        **opt_kwargs_dict,\n",
        "    )\n",
        "    # Some optimizers report param_norms and update_norms, so check if the\n",
        "    # returned value is a tuple which contains those or not.\n",
        "    model_opt_out = apply_grad_ret_list\n",
        "    # Note: uses validation inputs/targets and model_opt\n",
        "    with nn.stochastic(dropout_rng):\n",
        "      guide_logits = model_opt_out.target(\n",
        "          inputs=guide_inputs,\n",
        "          targets=guide_targets,\n",
        "          inputs_positions=guide_inputs_positions,\n",
        "          targets_positions=guide_targets_positions,\n",
        "          inputs_segmentation=guide_inputs_segmentation,\n",
        "          targets_segmentation=guide_targets_segmentation,\n",
        "          use_bfloat16=use_bfloat16,\n",
        "          dropout_rate=0,\n",
        "          attention_dropout_rate=0,\n",
        "          train=False,\n",
        "          cache=None)\n",
        "\n",
        "    guide_loss, guide_weight_sum = utils.compute_weighted_cross_entropy(\n",
        "        guide_logits,\n",
        "        guide_targets,\n",
        "        weights=jnp.where(guide_targets > 0, 1, 0).astype(jnp.float32))\n",
        "\n",
        "    total_guide_loss = guide_loss / guide_weight_sum\n",
        "    return total_guide_loss, (guide_logits, full_train_loss, train_logits,\n",
        "                              model_opt_out, model_learning_rate)\n",
        "\n",
        "  # Calculate gradient for all guided vars.\n",
        "  guided_hparam_names = [x for x in optimizer_dict.keys() if x != 'model'\n",
        "                        ]  # Exclude the model itself.\n",
        "  # Sort by the key (lp-<abbreviated param name>__<hparam> for param strings).\n",
        "  guided_hparam_names.sort()\n",
        "\n",
        "  # Extract values only.\n",
        "  guided_hparam_raw_values = [raw_vars_dict[k] for k in guided_hparam_names]\n",
        "\n",
        "  # This allows us to know which of the guided_hparam_raw_values corresponds to\n",
        "  # which name, and apply the appropriate logic, within the\n",
        "  # loss_fn_guidance_step.\n",
        "  map_guided_names_to_raw_indices = {\n",
        "      name: i for i, name in enumerate(guided_hparam_names)\n",
        "  }\n",
        "  # We have to use args list instead of kwargs dict for jax.grad to be able to\n",
        "  # pick out the correct arg for differentiation.\n",
        "  guide_step_args = [\n",
        "      optimizer_dict['model'],\n",
        "      map_guided_names_to_raw_indices,\n",
        "      *guided_hparam_raw_values,\n",
        "  ]\n",
        "\n",
        "  # The count of explicit args in the signature of loss_fn_guidance_step\n",
        "  # preceding the catch-all *per_param_raw_values.\n",
        "  # The specific mapping is:\n",
        "  # def loss_fn_guidance_step(\n",
        "  #   model_opt, -- index 0\n",
        "  #   map_guided_names_to_raw_indices, -- index 1\n",
        "  #   *per_param_raw_values,\n",
        "  #   )\n",
        "  # TODO(lichtarge): set this with the inspect.signature fn instead of manually.\n",
        "  max_index_explicit_args = 1\n",
        "  # Add 1 because enumerate starts at 0 and otherwise we'd overlap the index of\n",
        "  # map_guided_names_to_raw_indices.\n",
        "  guided_hparam_arg_indices = [\n",
        "      map_guided_names_to_raw_indices[x] + 1 + max_index_explicit_args\n",
        "      for x in guided_hparam_names\n",
        "  ]\n",
        "\n",
        "  # We do not know how many learned opts there will be so we unpack into a list.\n",
        "  # https://www.python.org/dev/peps/pep-3132/\n",
        "  # The unused vars are unregularized_loss and regularization_loss respectively.\n",
        "  (_, (guide_logits, _, train_logits, model_opt_out,\n",
        "       model_learning_rate)), *learned_grads = jax.value_and_grad(\n",
        "           loss_fn_guidance_step, guided_hparam_arg_indices,\n",
        "           has_aux=True)(*guide_step_args)\n",
        "\n",
        "  metrics_joint.update(\n",
        "      utils.compute_metrics(\n",
        "          train_logits,\n",
        "          targets,\n",
        "          weights=jnp.where(targets > 0, 1, 0).astype(jnp.float32),\n",
        "          tag='dataset_train'))\n",
        "  optimizer_dict['model'] = model_opt_out\n",
        "\n",
        "  # These are always used.\n",
        "  metrics_joint['hparams/guided_lr_schedule'] = model_learning_rate\n",
        "\n",
        "  if output_metrics is not None:\n",
        "    for guided_hparam in guided_hparam_names:\n",
        "      if guided_hparam == 'hp-learning_rate':\n",
        "        metrics_joint['guided_hparams/guided_lr_schedule'] = model_learning_rate\n",
        "      activated_vars = act_fn_dict[guided_hparam](raw_vars_dict[guided_hparam])\n",
        "      metrics_joint['guided_hparams/%s' % guided_hparam] = activated_vars\n",
        "\n",
        "  metrics_joint['hparams/default_lr_schedule'] = learning_rate_fn(current_step)\n",
        "\n",
        "  for i, guided_opt_key in enumerate(guided_hparam_names):\n",
        "    guided_optimizer = optimizer_dict[guided_opt_key]\n",
        "    # learned_grads is a List of size 1, the single element is a Tuple\n",
        "    raw_vars_grad = learned_grads[0][i]\n",
        "    if output_metrics is not None:\n",
        "      metrics_joint['guided_hparams_grads/%s' % guided_opt_key] = raw_vars_grad\n",
        "    if FLAGS.clip_nan_grads:\n",
        "      raw_vars_grad = jnp.nan_to_num(raw_vars_grad, posinf=10.0, neginf=-10.0)\n",
        "\n",
        "    # Deal with parallelization.\n",
        "    raw_vars_grad = jax.lax.pmean(raw_vars_grad, 'batch')\n",
        "    var_type = guided_parameters.get_var_type_from_opt_key(guided_opt_key)\n",
        "\n",
        "    # Grad clipping.\n",
        "    grad_clip_limit = FLAGS.grad_clip_limit\n",
        "    if grad_clip_limit != 0.0 and FLAGS.use_grad_clip:\n",
        "      metrics_joint['model_update/%s_grad_clip_limit' %\n",
        "                    guided_opt_key] = grad_clip_limit\n",
        "      metrics_joint['model_update/%s_num_grad_clipped_total' %\n",
        "                    guided_opt_key] = jnp.count_nonzero(\n",
        "                        raw_vars_grad > grad_clip_limit) + jnp.count_nonzero(\n",
        "                            raw_vars_grad < -grad_clip_limit)\n",
        "      raw_vars_grad = jax.tree_map(\n",
        "          lambda g_: jnp.clip(  # pylint: disable=g-long-lambda\n",
        "              g_,\n",
        "              a_min=-grad_clip_limit,\n",
        "              a_max=grad_clip_limit),\n",
        "          raw_vars_grad)\n",
        "\n",
        "    guided_update_kwargs = {'learning_rate': lr_guided_vars}\n",
        "    if guided_vars_dict[var_type]['optimizer_type'] == 'adam':\n",
        "      if FLAGS.guided_update_beta1 is not None:\n",
        "        guided_update_kwargs.update({'beta1': FLAGS.guided_update_beta1})\n",
        "      if FLAGS.guided_update_beta2 is not None:\n",
        "        guided_update_kwargs.update({'beta2': FLAGS.guided_update_beta2})\n",
        "      if FLAGS.guided_update_eps is not None:\n",
        "        guided_update_kwargs.update({'eps': FLAGS.guided_update_eps})\n",
        "      if FLAGS.guided_update_wd is not None:\n",
        "        guided_update_kwargs.update({'weight_decay': FLAGS.guided_update_wd})\n",
        "    elif guided_vars_dict[var_type]['optimizer_type'] in ['momentum', 'mom']:\n",
        "      if FLAGS.guided_update_beta1 is not None:\n",
        "        guided_update_kwargs.update({'beta': FLAGS.guided_update_beta1})\n",
        "      guided_update_kwargs.update({'nesterov': True})\n",
        "    elif guided_vars_dict[var_type]['optimizer_type'] in ['adagrad']:\n",
        "      if FLAGS.guided_update_eps is not None:\n",
        "        guided_update_kwargs.update({'eps': FLAGS.guided_update_eps})\n",
        "    elif guided_vars_dict[var_type]['optimizer_type'] in ['adafactor']:\n",
        "      if FLAGS.guided_update_eps is not None:\n",
        "        guided_update_kwargs.update({'eps': FLAGS.guided_update_eps})\n",
        "      if FLAGS.guided_update_beta1 is not None:\n",
        "        guided_update_kwargs.update({'beta1': FLAGS.guided_update_beta1})\n",
        "    # apply raw vars gradient to dp optimizer, updating raw_vars\n",
        "\n",
        "    guided_optimizer = guided_optimizer.apply_gradient(raw_vars_grad,\n",
        "                                                       **guided_update_kwargs)\n",
        "    # update the optimizer_dict\n",
        "    optimizer_dict[guided_opt_key] = guided_optimizer\n",
        "  if output_metrics is not None:\n",
        "    metrics_joint.update(\n",
        "        utils.compute_metrics(\n",
        "            guide_logits,\n",
        "            guide_targets,\n",
        "            weights=jnp.where(guide_targets > 0, 1, 0).astype(jnp.float32),\n",
        "            tag='dataset_guide'))\n",
        "\n",
        "  return optimizer_dict, metrics_joint, new_dropout_rng\n",
        "\n"
      ],
      "metadata": {
        "id": "fGEeP7nBXBBB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Training Code\n",
        "\n",
        "\n",
        "def _get_default_opt_kwargs(model_optimizer_type, lr):\n",
        "  \"\"\"Returns a dict of all default hparams of the given optimizer.\"\"\"\n",
        "  opt_kwargs = {'learning_rate': lr}\n",
        "  model_optimizer_type = model_optimizer_type.split('_')\n",
        "  if 'adam' in model_optimizer_type:\n",
        "    opt_kwargs.update({\n",
        "        'beta1': FLAGS.adam_beta1,\n",
        "        'beta2': FLAGS.adam_beta2,\n",
        "        'eps': FLAGS.adam_eps,\n",
        "        'weight_decay': FLAGS.adam_weight_decay,\n",
        "    })\n",
        "  elif 'lamb' in model_optimizer_type:\n",
        "    opt_kwargs.update({\n",
        "        'beta1': FLAGS.lamb_beta1,\n",
        "        'beta2': FLAGS.lamb_beta2,\n",
        "        'eps': FLAGS.lamb_eps,\n",
        "        'weight_decay': FLAGS.lamb_weight_decay,\n",
        "    })\n",
        "  elif 'mom' in model_optimizer_type:\n",
        "    opt_kwargs.update({\n",
        "        'beta': FLAGS.momentum_beta1,\n",
        "        'weight_decay': FLAGS.momentum_weight_decay,\n",
        "        'nesterov': FLAGS.momentum_nesterov,\n",
        "    })\n",
        "  elif 'adagrad' in model_optimizer_type:\n",
        "    opt_kwargs.update({\n",
        "        'eps': FLAGS.adagrad_eps,\n",
        "    })\n",
        "  elif 'adafactor' in model_optimizer_type:\n",
        "    opt_kwargs.update({\n",
        "        'beta1':\n",
        "            FLAGS.adafactor_beta1,\n",
        "        'decay_rate':\n",
        "            FLAGS.adafactor_decay_rate,\n",
        "        'epsilon1':\n",
        "            FLAGS.adafactor_epsilon1,\n",
        "        'epsilon2':\n",
        "            FLAGS.adafactor_epsilon2,\n",
        "        'factored':\n",
        "            FLAGS.adafactor_factored,\n",
        "        'multiply_by_parameter_scale':\n",
        "            FLAGS.adafactor_multiply_by_parameter_scale,\n",
        "        'step_offset':\n",
        "            FLAGS.adafactor_step_offset,\n",
        "        'clipping_threshold':\n",
        "            FLAGS.adafactor_clipping_threshold,\n",
        "        'weight_decay_rate':\n",
        "            FLAGS.adafactor_weight_decay_rate,\n",
        "        'min_dim_size_to_factor':\n",
        "            FLAGS.adafactor_min_dim_size_to_factor,\n",
        "    })\n",
        "  elif 'sgd' in model_optimizer_type or 'gd' in model_optimizer_type:\n",
        "    pass\n",
        "  else:\n",
        "    raise ValueError('Unrecognized model_optimizer_type: %s' %\n",
        "                     model_optimizer_type)\n",
        "  return opt_kwargs\n",
        "\n",
        "\n",
        "def get_step_from_opt(model_optimizer):\n",
        "  if not hasattr(model_optimizer, 'state'):\n",
        "    raise ValueError(\n",
        "        'model optimizer does not have state attribute! Cannot retrieve step.')\n",
        "  if isinstance(model_optimizer.state, tuple):\n",
        "    if hasattr(model_optimizer.state[0], 'count'):\n",
        "      return model_optimizer.state[0].count\n",
        "  if hasattr(model_optimizer.state, 'count'):\n",
        "    return model_optimizer.state.count\n",
        "  return model_optimizer.state.step\n",
        "\n",
        "\n",
        "def update_model_step(\n",
        "    optimizer_dict,\n",
        "    train_batch,\n",
        "    learning_rate_fn,\n",
        "    use_bfloat16=False,\n",
        "    dropout_rng=None,\n",
        "    output_metrics=None,\n",
        "):\n",
        "  \"\"\"Updates the model only, without using guided parameters.\n",
        "\n",
        "  Args:\n",
        "    optimizer_dict: a dictionary holding all optimizers\n",
        "    train_batch: the batch of training data to be applied\n",
        "    learning_rate_fn: model learning rate function\n",
        "    use_bfloat16: bool, If True, use lower-precision bfloat16\n",
        "    dropout_rng: the JAX rng for dropout\n",
        "    output_metrics: metrics to be passed through and reported\n",
        "\n",
        "  Returns:\n",
        "    optimizer_dict: updated with steps taken for guided parameter optimizers\n",
        "    metrics: metrics dict keeping track of tensorboard-reported summary values\n",
        "    new_dropout_rng: updated RNG\n",
        "  \"\"\"\n",
        "  metrics_train = {}  # for summarizing\n",
        "  if output_metrics:\n",
        "    metrics_train.update(output_metrics)\n",
        "  model_optimizer = optimizer_dict['model']  # for readability\n",
        "  dropout_rng, new_dropout_rng = random.split(dropout_rng)\n",
        "  current_step = model_optimizer.state.step\n",
        "\n",
        "  # Extract fields of interest from the training batch.\n",
        "  (inputs, targets, inputs_positions, targets_positions, inputs_segmentation,\n",
        "   targets_segmentation) = [train_batch.get(k, None) for k in _TRAIN_KEYS_BASE]\n",
        "\n",
        "  # Loss fn for model step lookahead, used within loss_fn_guidance_step.\n",
        "  # Applies example weighting, smoothing, etc, to the lookahead model.\n",
        "  def loss_fn_model_lookahead(model,):\n",
        "    \"\"\"Loss fn for model step lookahead, used within loss_fn_guidance_step.\"\"\"\n",
        "\n",
        "    # Default to FLAGS value if not set within guided_vars_dict.\n",
        "    dropout_rate = FLAGS.dropout_rate\n",
        "    attention_dropout_rate = FLAGS.attention_dropout_rate\n",
        "    label_smoothing = FLAGS.label_smoothing if FLAGS.label_smoothing else 0.0\n",
        "\n",
        "    with nn.stochastic(dropout_rng):\n",
        "      train_logits = model(\n",
        "          inputs=inputs,\n",
        "          targets=targets,\n",
        "          use_bfloat16=use_bfloat16,\n",
        "          inputs_positions=inputs_positions,\n",
        "          targets_positions=targets_positions,\n",
        "          inputs_segmentation=inputs_segmentation,\n",
        "          targets_segmentation=targets_segmentation,\n",
        "          train=True,\n",
        "          cache=None,\n",
        "          dropout_rate=dropout_rate,\n",
        "          attention_dropout_rate=attention_dropout_rate)\n",
        "\n",
        "    # Ignore loss on padding.\n",
        "    weights = jnp.where(targets > 0, 1, 0).astype(jnp.float32)\n",
        "\n",
        "    loss, weight_sum = utils.compute_weighted_cross_entropy(\n",
        "        train_logits,\n",
        "        targets,\n",
        "        weights,\n",
        "        label_smoothing,\n",
        "    )\n",
        "    full_train_loss = loss / jnp.where(weight_sum != 0, weight_sum, 1)\n",
        "\n",
        "    return full_train_loss, (train_logits,)\n",
        "\n",
        "  model_learning_rate = learning_rate_fn(current_step)\n",
        "  (_, (train_logits,)), model_train_step_grad = jax.value_and_grad(\n",
        "      loss_fn_model_lookahead, argnums=0, has_aux=True)(model_optimizer.target,)\n",
        "  model_train_step_grad = jax.lax.pmean(model_train_step_grad, 'batch')\n",
        "\n",
        "  opt_kwargs_dict = _get_default_opt_kwargs(\n",
        "      model_optimizer_type=FLAGS.model_optimizer_type,\n",
        "      lr=model_learning_rate,\n",
        "  )\n",
        "  model_opt_out = model_optimizer.apply_gradient(\n",
        "      model_train_step_grad,\n",
        "      **opt_kwargs_dict,\n",
        "  )\n",
        "\n",
        "  metrics_train.update(\n",
        "      utils.compute_metrics(\n",
        "          train_logits,\n",
        "          targets,\n",
        "          weights=jnp.where(targets > 0, 1, 0).astype(jnp.float32),\n",
        "          tag='dataset_train'))\n",
        "  optimizer_dict['model'] = model_opt_out\n",
        "\n",
        "  # These are always used.\n",
        "  metrics_train['hparams/default_lr_schedule'] = learning_rate_fn(current_step)\n",
        "\n",
        "  return optimizer_dict, metrics_train, new_dropout_rng\n",
        "\n",
        "\n",
        "def eval_step(model, batch, label_smoothing=0.0, use_bfloat16=False):\n",
        "  \"\"\"Calculate evaluation metrics on a batch.\"\"\"\n",
        "  del label_smoothing\n",
        "  (inputs, targets, inputs_positions, targets_positions, inputs_segmentation,\n",
        "   targets_segmentation) = [batch.get(k, None) for k in _TRAIN_KEYS_BASE]\n",
        "  weights = jnp.where(targets > 0, 1.0, 0.0)\n",
        "  logits = model(\n",
        "      inputs,\n",
        "      targets,\n",
        "      use_bfloat16=use_bfloat16,\n",
        "      inputs_positions=inputs_positions,\n",
        "      targets_positions=targets_positions,\n",
        "      inputs_segmentation=inputs_segmentation,\n",
        "      targets_segmentation=targets_segmentation,\n",
        "      train=False,\n",
        "      cache=None)\n",
        "  return utils.compute_metrics(logits, targets, weights, tag='dataset_eval')\n",
        "\n",
        "\n",
        "def batch_to_numpy(batch):\n",
        "  return jax.tree_map(lambda x: x._numpy(), batch)  # pylint: disable=protected-access\n",
        "\n",
        "\n",
        "def do_eval_single(\n",
        "    step,\n",
        "    eval_ds,\n",
        "    eval_ds_name,\n",
        "    p_eval_step,\n",
        "    optimizer,\n",
        "):\n",
        "  \"\"\"Computes metrics on a single evaluation dataset.\"\"\"\n",
        "  eval_metrics = []\n",
        "  eval_iter = iter(eval_ds)\n",
        "  sum_unique_examples = 0\n",
        "  logging.info('Starting eval for eval set %s at step %s.', eval_ds_name, step)\n",
        "  if FLAGS.print_for_colab:\n",
        "    print('Starting eval for eval set %s at step %s.', eval_ds_name, step)\n",
        "\n",
        "  eval_ds_num_batches = 0\n",
        "  for eval_batch in eval_iter:\n",
        "    eval_ds_num_batches += 1\n",
        "    eval_batch = batch_to_numpy(eval_batch)\n",
        "    sum_unique_examples += data.get_unique_examples(eval_batch)\n",
        "    eval_batch = common_utils.shard(eval_batch)\n",
        "    metrics = p_eval_step(optimizer.target, eval_batch)\n",
        "    eval_metrics.append(metrics)\n",
        "  if not eval_metrics:\n",
        "    raise ValueError('Eval has failed for: %s' % eval_ds_name)\n",
        "  eval_metrics = common_utils.get_metrics(eval_metrics)\n",
        "  eval_metrics_sums = jax.tree_map(jnp.sum, eval_metrics)\n",
        "\n",
        "  eval_denominator = eval_metrics_sums.pop('dataset_eval/denominator')\n",
        "  eval_summary = jax.tree_map(\n",
        "      lambda x: x / eval_denominator,  # pylint: disable=cell-var-from-loop\n",
        "      eval_metrics_sums)\n",
        "  eval_summary['dataset_eval/unique_examples'] = sum_unique_examples\n",
        "\n",
        "  logging.info('Finished do_eval (%s batches).', eval_ds_num_batches)\n",
        "  if FLAGS.print_for_colab:\n",
        "    print('Finished do_eval (%s batches).', eval_ds_num_batches)\n",
        "  return eval_summary\n",
        "\n",
        "\n",
        "def write_finished_training(model_dir):\n",
        "  \"\"\"Write a file into CNS marking this training run as complete..\"\"\"\n",
        "  if jax.process_index() != 0:\n",
        "    return\n",
        "  logging.info('Writing to CNS that training is done.')\n",
        "  try:\n",
        "    with tf.io.gfile.GFile(\n",
        "        os.path.join(model_dir, _TRAINING_DONE_FILENAME), 'w') as f:\n",
        "      f.write('Training run is complete.')\n",
        "  except tf.errors.NotFoundError:\n",
        "    pass\n",
        "\n",
        "\n",
        "def maybe_remove_finished_training(model_dir):\n",
        "  \"\"\"Remove (if present) the CNS file marking this training run as complete.\"\"\"\n",
        "  if jax.process_index() != 0:\n",
        "    return\n",
        "  if tf.io.gfile.exists(os.path.join(model_dir, _TRAINING_DONE_FILENAME)):\n",
        "    logging.info('Removing \"%s\" while current run ongoing.',\n",
        "                 _TRAINING_DONE_FILENAME)\n",
        "    tf.io.gfile.remove(os.path.join(model_dir, _TRAINING_DONE_FILENAME))\n",
        "\n",
        "\n",
        "def save_all_checkpoints(\n",
        "    t_start,\n",
        "    optimizer_dict,\n",
        "    step,\n",
        "):\n",
        "  \"\"\"Save optimizers from optimizer_dict.\"\"\"\n",
        "  if not FLAGS.save_checkpoints:\n",
        "    return\n",
        "  # only save once, not once per worker\n",
        "  if jax.process_index() == 0:\n",
        "    utils.log_message(t_start, 'Saving optimizers, current step: %s' % step)\n",
        "    if FLAGS.print_for_colab:\n",
        "      print('saving optimizers, current step: %s' % step)\n",
        "    print('sorted(list(optimizer_dict.keys()))')\n",
        "    print(sorted(list(optimizer_dict.keys())))\n",
        "    for opt_dict_key in sorted(optimizer_dict.keys()):\n",
        "      optimizer = optimizer_dict[opt_dict_key]\n",
        "      utils.log_message(t_start, 'Saving optimizer: %s' % opt_dict_key)\n",
        "      prefix_string = '%s_' % opt_dict_key.replace('/', '.')\n",
        "      write_to_file = os.path.join(FLAGS.model_dir,\n",
        "                                   '%s%s' % (prefix_string, step))\n",
        "      # It may be necessary to overwrite hp-vars files if the training run\n",
        "      # was interrupted and we had to go back to the last model ckpt.\n",
        "      if tf.io.gfile.exists(write_to_file):\n",
        "        logging.info('Skipping saving ckpt at %s because it already exists.',\n",
        "                     write_to_file)\n",
        "        print('Skipping saving ckpt at %s because it already exists.',\n",
        "              write_to_file)\n",
        "        continue\n",
        "      checkpoints.save_checkpoint(\n",
        "          FLAGS.model_dir,\n",
        "          jax_utils.unreplicate(optimizer),\n",
        "          step,\n",
        "          prefix=prefix_string,\n",
        "          keep=FLAGS.keep_checkpoints_count)\n",
        "\n",
        "\n",
        "def maybe_save_ckpts(\n",
        "    t_start,\n",
        "    optimizer_dict,\n",
        "    current_step,\n",
        "    steps_since_last_model_ckpt,\n",
        "    steps_since_last_guided_ckpt,\n",
        "):\n",
        "  \"\"\"Save optimizers from optimizer_dict.\"\"\"\n",
        "  if not FLAGS.save_checkpoints:\n",
        "    return steps_since_last_model_ckpt, steps_since_last_guided_ckpt\n",
        "  # only save once, not once per worker\n",
        "  if jax.process_index() != 0:\n",
        "    return steps_since_last_model_ckpt, steps_since_last_guided_ckpt\n",
        "  # Save a checkpoint on one host after every model_ckpt_min_freq steps.\n",
        "  # Maybe save guided vars more frequently (every guided_ckpt_min_freq steps).\n",
        "  save_guided = False\n",
        "  save_model = False\n",
        "\n",
        "  if FLAGS.guided_ckpt_min_freq and steps_since_last_guided_ckpt >= FLAGS.guided_ckpt_min_freq:\n",
        "    save_guided = True\n",
        "    steps_since_last_guided_ckpt = 0\n",
        "  else:\n",
        "    steps_since_last_guided_ckpt += 1\n",
        "  if FLAGS.model_ckpt_min_freq and steps_since_last_model_ckpt >= FLAGS.model_ckpt_min_freq - 1:\n",
        "    save_model = True\n",
        "    steps_since_last_model_ckpt = 0\n",
        "    save_guided = True\n",
        "    steps_since_last_guided_ckpt = 0\n",
        "  else:\n",
        "    steps_since_last_model_ckpt += 1\n",
        "\n",
        "  for opt_dict_key, optimizer in optimizer_dict.items():\n",
        "    if opt_dict_key == 'model' and not save_model:\n",
        "      continue\n",
        "    if opt_dict_key != 'model' and not save_guided:\n",
        "      continue\n",
        "    utils.log_message(t_start, 'saving optimizer: %s' % opt_dict_key)\n",
        "    prefix_string = '%s_' % opt_dict_key\n",
        "    write_to_file = os.path.join(FLAGS.model_dir,\n",
        "                                 '%s%s' % (prefix_string, current_step))\n",
        "    # it may be necessary to overwrite guided-vars files if the training run\n",
        "    # was interrupted and we had to go back to the last model ckpt\n",
        "    overwrite = False\n",
        "    if tf.io.gfile.exists(write_to_file):\n",
        "      if opt_dict_key == 'model':\n",
        "        raise ValueError('trying to save model but already exists... '\n",
        "                         'this can happen if two runs are writing to same dir?'\n",
        "                         ' file: %s' % write_to_file)\n",
        "      utils.log_message(\n",
        "          t_start, 'Overwriting old guided vars checkpoint: %s' % write_to_file)\n",
        "      tf.io.gfile.remove(write_to_file)\n",
        "      overwrite = True\n",
        "    checkpoints.save_checkpoint(\n",
        "        FLAGS.model_dir,\n",
        "        jax_utils.unreplicate(optimizer),\n",
        "        current_step,\n",
        "        prefix=prefix_string,\n",
        "        overwrite=overwrite,\n",
        "        keep=FLAGS.keep_checkpoints_count)\n",
        "  return steps_since_last_model_ckpt, steps_since_last_guided_ckpt\n",
        "\n",
        "\n",
        "def write_flags_files(output_dir, guided_vars_dict):\n",
        "  \"\"\"Write run flags to a human-readable text file and a json dump.\"\"\"\n",
        "  if jax.process_index() != 0:\n",
        "    return\n",
        "\n",
        "  if not tf.io.gfile.exists(output_dir):\n",
        "    tf.io.gfile.MakeDirs(output_dir)\n",
        "\n",
        "  flag_dict = FLAGS.flags_by_module_dict()\n",
        "  jax_module_names = [x for x in flag_dict.keys() if 'jax' in x or 'train' in x]\n",
        "\n",
        "  jax_flags_str = ''\n",
        "  jax_train_dict = {}\n",
        "  for module in jax_module_names:\n",
        "    jax_flags_str += '\\n' + module + '\\n'\n",
        "    jax_flags_str += '\\t\\n'.join(\n",
        "        sorted(['--%s=%s' % (f.name, f.value) for f in flag_dict[module]]))\n",
        "    jax_flags_str += '\\n'\n",
        "\n",
        "    # This isolates train.py flags locally and on borg.\n",
        "    if 'jax.train' in module or 'train.par' in module:\n",
        "      jax_train_dict.update({f.name: f.value for f in flag_dict[module]})\n",
        "  flags_txt = os.path.join(output_dir, 'jax_flags.txt')\n",
        "  logging.info('JAX host: %d / %d', jax.process_index(), jax.process_count())\n",
        "\n",
        "  logging.info('writing flags into: %s', flags_txt)\n",
        "  try:\n",
        "    with tf.io.gfile.GFile(flags_txt, 'w') as f:\n",
        "      f.write(jax_flags_str)\n",
        "  except tf.errors.NotFoundError:\n",
        "    logging.warn('Failed to write flags to CNS: %s', flags_txt)\n",
        "\n",
        "  guided_flags_json = os.path.join(output_dir, 'train_run_hparams.json')\n",
        "  if guided_vars_dict:\n",
        "    flattened_dict = {}\n",
        "    for gv_key, gv_dict in guided_vars_dict.items():\n",
        "      flattened_dict.update({\n",
        "          '%s.%s' % (gv_key, k): v\n",
        "          for k, v in gv_dict.items()\n",
        "          if k != 'raw_guided_vars'\n",
        "      })  # Exclude raw_guided_vars as it is unnecessary, non-JSON-serializable.\n",
        "\n",
        "    jax_train_dict.update(\n",
        "        {'guided_vars.%s' % k: v for k, v in flattened_dict.items()})\n",
        "    jax_train_dict = guided_parameters.make_dict_json_safe(jax_train_dict)\n",
        "  try:\n",
        "    with tf.io.gfile.GFile(guided_flags_json, 'w') as f:\n",
        "      json.dump(jax_train_dict, f)\n",
        "  except tf.errors.NotFoundError:\n",
        "    pass\n",
        "\n",
        "\n",
        "def _load_spm_tokenizer(model_path):\n",
        "  \"\"\"Load a spm tokenizer from given model filepath.\"\"\"\n",
        "  with tf.io.gfile.GFile(model_path, 'rb') as f:\n",
        "    spm_model = f.read()\n",
        "  sp_tokenizer = spm.SentencePieceProcessor()\n",
        "  sp_tokenizer.LoadFromSerializedProto(spm_model)\n",
        "  return sp_tokenizer\n",
        "\n",
        "\n",
        "def restore_or_init_model(t_start):\n",
        "  \"\"\"Initializes a model; restores from prior ckpt or init ckpt as needed.\n",
        "\n",
        "  Randomly initialize a Transformer model with FLAGS-specified parameters.\n",
        "  If the current job has been pre-empted (or started again after termination),\n",
        "    restore the model parameters from the most recent checkpoint in the output\n",
        "    directory.\n",
        "  If no output dir ckpts exist, initialize from the FLAGS.init_checkpoint\n",
        "    parameters.\n",
        "\n",
        "  Args:\n",
        "    t_start: timestamp used for logging\n",
        "\n",
        "  Returns:\n",
        "    model_optimizer: the model optimizer\n",
        "    optimizer_step: the step of the current optimizer state\n",
        "    rng: jax PRNG key, an np.ndarray of two ints\n",
        "  \"\"\"\n",
        "  utils.log_message(t_start, 'Initializing model.')\n",
        "  encoder = _load_spm_tokenizer(FLAGS.vocab_path)\n",
        "  vocab_size = int(encoder.GetPieceSize())\n",
        "  current_run_step = 0\n",
        "\n",
        "  # Build Model and Optimizer\n",
        "  transformer_kwargs = {\n",
        "      'vocab_size': vocab_size,\n",
        "      'output_vocab_size': vocab_size,\n",
        "      'emb_dim': FLAGS.emb_dim,\n",
        "      'num_heads': FLAGS.num_heads,\n",
        "      'num_layers': FLAGS.num_layers,\n",
        "      'qkv_dim': FLAGS.qkv_dim,\n",
        "      'mlp_dim': FLAGS.mlp_dim,\n",
        "      'max_len': FLAGS.max_target_length,\n",
        "      'share_embeddings': FLAGS.share_embeddings,\n",
        "      'logits_via_embedding': FLAGS.logits_via_embedding,\n",
        "  }\n",
        "  logging.info('Transformer kwargs: %s', transformer_kwargs)\n",
        "\n",
        "  rng = random.PRNGKey(FLAGS.jax_random_seed)\n",
        "  rng, init_rng = random.split(rng)\n",
        "\n",
        "  input_shape = (FLAGS.batch_size, FLAGS.max_target_length)\n",
        "  target_shape = (FLAGS.batch_size, FLAGS.max_target_length)\n",
        "  model, cache_def = models.create_model(init_rng, input_shape, target_shape,\n",
        "                                         transformer_kwargs)\n",
        "  del cache_def\n",
        "  optimizer = models.init_optimizer_by_type(model, FLAGS.model_optimizer_type)\n",
        "  # We access model only from optimizer below via optimizer.target.\n",
        "  del model\n",
        "  model = None\n",
        "\n",
        "  # Remove bad tmp files caused by preemption during save; avoids error on save.\n",
        "  if jax.process_index() == 0:\n",
        "    for x in tf.io.gfile.glob(os.path.join(FLAGS.model_dir, '*_tmp')):\n",
        "      tf.io.gfile.remove(x)\n",
        "  glob_path = tf.io.gfile.glob(os.path.join(FLAGS.model_dir, 'model_*'))\n",
        "  if glob_path:\n",
        "    # If training has already started, and checkpoints exist in the model_dir,\n",
        "    # take the last checkpoint available (occurs if training job is preempted).\n",
        "    max_model_step = int(\n",
        "        checkpoints.natural_sort(glob_path)[-1].split('model_')[-1])\n",
        "    current_run_step = max_model_step\n",
        "    # restore_checkpoint checks model_dir for present checkpoints;\n",
        "    # if no checkpoints in model_dir, it passes the optimizer unchanged\n",
        "    optimizer = checkpoints.restore_checkpoint(\n",
        "        FLAGS.model_dir, optimizer, prefix='model_')\n",
        "\n",
        "    # Grab last step from the optimizer itself.\n",
        "    # this is progress already made this training run [0 --> num_train_steps]\n",
        "    optimizer_step = get_step_from_opt(optimizer)\n",
        "    utils.log_message(\n",
        "        t_start, 'Restoring ckpt from \"model_%s\", optimizer step (including '\n",
        "        'warm-initialization steps): %s, current run step (excluding warm-init)'\n",
        "        ': %s' % (max_model_step, optimizer_step, current_run_step))\n",
        "    if FLAGS.print_for_colab:\n",
        "      print(\n",
        "          t_start, 'Restoring ckpt from \"model_%s\", optimizer step (including '\n",
        "          'warm-initialization steps): %s, current run step (excluding warm-init)'\n",
        "          ': %s' % (max_model_step, optimizer_step, current_run_step))\n",
        "\n",
        "  # if training hasn't already started / saved a ckpt, init from init_checkpoint\n",
        "  elif FLAGS.init_checkpoint:\n",
        "    utils.log_message(t_start,\n",
        "                      'Init checkpoint from %s' % FLAGS.init_checkpoint)\n",
        "    if FLAGS.print_for_colab:\n",
        "      print(t_start, 'Init checkpoint from %s' % FLAGS.init_checkpoint)\n",
        "    with tf.io.gfile.GFile(FLAGS.init_checkpoint, 'rb') as fp:\n",
        "      optimizer = serialization.from_bytes(optimizer, fp.read())\n",
        "    optimizer_step = get_step_from_opt(optimizer)\n",
        "    if FLAGS.take_current_run_step_from_init:\n",
        "      current_run_step = optimizer_step\n",
        "\n",
        "  return optimizer, current_run_step, rng, model\n",
        "\n",
        "\n",
        "def remove_stale_guided_vars_ckpts(opt_key, restore_step):\n",
        "  \"\"\"Remove any stale checkpoints which may have been left by task failure.\"\"\"\n",
        "  glob_paths = tf.io.gfile.glob(os.path.join(FLAGS.model_dir, '%s*' % opt_key))\n",
        "  past_restore_path = os.path.join(FLAGS.model_dir,\n",
        "                                   '%s_%s' % (opt_key, restore_step + 1))\n",
        "  glob_paths.append(past_restore_path)\n",
        "  # It may be necessary to overwrite guided optimizer ckpts if the training run\n",
        "  # was interrupted and we had to go back to the last model ckpt.\n",
        "  sorted_ckpts = checkpoints.natural_sort(glob_paths)\n",
        "  for x in sorted_ckpts[sorted_ckpts.index(past_restore_path):]:\n",
        "    if tf.io.gfile.exists(x):\n",
        "      logging.info(\n",
        "          'Removing bad ckpt (got ahead of model due to task failure): %s', x)\n",
        "      tf.io.gfile.remove(x)\n",
        "\n",
        "\n",
        "def restore_or_init_guided_optimizers(\n",
        "    t_start,\n",
        "    guided_vars_dict,\n",
        "):\n",
        "  \"\"\"Initializes and maybe restores guided optimizers.\n",
        "\n",
        "  Args:\n",
        "    t_start: timestamp used for logging\n",
        "    guided_vars_dict: the guided parameters to be applied in this run\n",
        "\n",
        "  Raises:\n",
        "    ValueError: if the guided_vars_dict contains an unrecognized key\n",
        "    ValueError: if the guided opt step is not equal to the expected restore step\n",
        "\n",
        "  Returns:\n",
        "    optimizer_dict: contains all guided optimizers\n",
        "\n",
        "  \"\"\"\n",
        "  optimizer_dict = {}\n",
        "  try:\n",
        "    if jax.process_index() == 0:\n",
        "      for x in tf.io.gfile.glob(os.path.join(FLAGS.model_dir, '*_tmp')):\n",
        "        tf.io.gfile.remove(x)\n",
        "    glob_path = tf.io.gfile.glob(os.path.join(FLAGS.model_dir, 'model_*'))\n",
        "  except tf.errors.NotFoundError:\n",
        "    glob_path = None\n",
        "  # If there are no models_* files, restore step is 0\n",
        "  restore_step = -1\n",
        "  if glob_path:\n",
        "    last_ckpt = checkpoints.natural_sort(glob_path)[-1]\n",
        "    if 'model_' in last_ckpt:\n",
        "      last_ckpt_num = last_ckpt.split('model_')[-1]\n",
        "      restore_step = int(last_ckpt_num)\n",
        "  if FLAGS.print_for_colab:\n",
        "    print('guided_var_types:')\n",
        "    for guided_var_type in guided_vars_dict:\n",
        "      print('\\t', guided_var_type)\n",
        "  for guided_var_type in guided_vars_dict:\n",
        "    guided_var_opt_key = guided_parameters.get_opt_key_from_var_type(\n",
        "        guided_var_type)\n",
        "    # First, remove stale checkpoints which may have been left behind by\n",
        "    # pre-emption.\n",
        "    remove_stale_guided_vars_ckpts(guided_var_opt_key, restore_step)\n",
        "    optimizer_hparams = {\n",
        "        'weight_decay': FLAGS.guided_weight_decay,\n",
        "    }\n",
        "    cur_guided_opt = models.init_optimizer_by_type(\n",
        "        guided_vars_dict[guided_var_type]['raw_guided_vars'],\n",
        "        guided_vars_dict[guided_var_type]['optimizer_type'],\n",
        "        optimizer_hparams=optimizer_hparams)\n",
        "\n",
        "    if restore_step > 0:\n",
        "      utils.log_message(\n",
        "          t_start, 'Restoring cur_guided_opt from \"%s_%s\", ' %\n",
        "          (guided_var_opt_key, restore_step))\n",
        "\n",
        "      cur_guided_opt = checkpoints.restore_checkpoint(\n",
        "          FLAGS.model_dir,\n",
        "          cur_guided_opt,\n",
        "          prefix='%s_%s' % (guided_var_opt_key, restore_step))\n",
        "      cur_opt_step = get_step_from_opt(cur_guided_opt)\n",
        "      utils.log_message(\n",
        "          t_start, 'restoring cur_guided_opt from \"hp-%s_%s\", '\n",
        "          'optimizer step: %s' % (guided_var_type, restore_step, cur_opt_step))\n",
        "      if FLAGS.print_for_colab:\n",
        "        print('restoring cur_guided_opt from \"hp-%s_%s\", '\n",
        "              'optimizer step: %s' %\n",
        "              (guided_var_type, restore_step, cur_opt_step))\n",
        "    elif 'init_path' in guided_vars_dict[guided_var_type]:\n",
        "      init_checkpoint_path = guided_vars_dict[guided_var_type]['init_path']\n",
        "      utils.log_message(\n",
        "          t_start,\n",
        "          'initializing cur_guided_opt from \"%s\", ' % init_checkpoint_path)\n",
        "      if FLAGS.print_for_colab:\n",
        "        print('initializing cur_guided_opt from \"%s\", ' % init_checkpoint_path)\n",
        "      init_dir, dp_init_step = init_checkpoint_path.split('%s_' %\n",
        "                                                          guided_var_opt_key)\n",
        "      if FLAGS.print_for_colab:\n",
        "        print('init_dir, dp_init_step')\n",
        "        print(init_dir, dp_init_step)\n",
        "      cur_guided_opt = checkpoints.restore_checkpoint(\n",
        "          init_dir,\n",
        "          cur_guided_opt,\n",
        "          prefix='%s_%s' % (guided_var_opt_key, dp_init_step))\n",
        "    cur_guided_opt = jax_utils.replicate(cur_guided_opt)\n",
        "    utils.log_message(\n",
        "        t_start, '%s optimizer initialized. optimizer step: %s' %\n",
        "        (guided_var_opt_key, get_step_from_opt(cur_guided_opt)))\n",
        "    optimizer_dict[guided_var_opt_key] = cur_guided_opt\n",
        "\n",
        "  return optimizer_dict\n",
        "\n",
        "\n",
        "def set_up_summary_writers():\n",
        "  \"\"\"Helper fn to set up metric-reporting SummaryWriters.\"\"\"\n",
        "  train_sum_writer = tensorboard.SummaryWriter(\n",
        "      os.path.join(FLAGS.model_dir, 'train-%s' % FLAGS.training_dataset))\n",
        "\n",
        "  eval_sum_writers_dict = {}\n",
        "  # TODO(lichtarge): fix this\n",
        "  for eval_ds_name in get_evalsets():\n",
        "    if FLAGS.print_for_colab:\n",
        "      print('running eval for: %s' % eval_ds_name)\n",
        "    tb_summary_name = 'eval-%s' % eval_ds_name\n",
        "    eval_sum_writers_dict[tb_summary_name] = tensorboard.SummaryWriter(\n",
        "        os.path.join(FLAGS.model_dir, tb_summary_name))\n",
        "\n",
        "  return train_sum_writer, eval_sum_writers_dict\n",
        "\n",
        "\n",
        "def get_dataset_from_path(data_path, batch_size):\n",
        "  \"\"\"Helper to fetch training dataset, specified by FLAGS.training_dataset.\n",
        "\n",
        "  Args:\n",
        "    data_path: path to data\n",
        "    batch_size: size of batch\n",
        "\n",
        "  Returns:\n",
        "    tf.data.Dataset\n",
        "  \"\"\"\n",
        "  train_ds_kwargs = {\n",
        "      'batch_size': batch_size,\n",
        "      'deterministic': FLAGS.print_for_colab,\n",
        "      'max_length': FLAGS.max_target_length,\n",
        "      'pack': True,\n",
        "      'random_seed': FLAGS.jax_random_seed,\n",
        "      'repeat': True,\n",
        "      'vocab_path': FLAGS.vocab_path,\n",
        "  }\n",
        "  logging.info('Train args: %s', train_ds_kwargs)\n",
        "  if FLAGS.print_for_colab:\n",
        "    print('Training data is deterministic.')\n",
        "    print('Train dataset args: %s' % train_ds_kwargs)\n",
        "  else:\n",
        "    logging.info('training is not deterministic')\n",
        "  train_ds = data.get_prepacked_examples(\n",
        "      file_pattern=data_path, **train_ds_kwargs)\n",
        "  return train_ds\n",
        "\n",
        "\n",
        "def maybe_get_new_guide_batch(guide_batch, guide_iter, frequency,\n",
        "                              guide_batch_unq_ex):\n",
        "  \"\"\"Update guide batch if FLAGS.guide_batch_update_freq == frequency.\"\"\"\n",
        "  if frequency not in ['PER_EPOCH', 'PER_TRAINING_STEP', 'PER_GUIDANCE_STEP']:\n",
        "    raise ValueError('Invalid frequency: %s' % frequency)\n",
        "  if FLAGS.guide_batch_update_freq == frequency:\n",
        "    next_guide_batch = batch_to_numpy(next(guide_iter))\n",
        "    guide_batch_unq_ex = data.get_unique_examples(next_guide_batch)\n",
        "    logging.info('guide batch unq ex: %s', guide_batch_unq_ex)\n",
        "    if FLAGS.print_for_colab:\n",
        "      print('guide batch unq ex:', guide_batch_unq_ex)\n",
        "    guide_batch = common_utils.shard(next_guide_batch)\n",
        "  return guide_batch, guide_iter, guide_batch_unq_ex\n",
        "\n",
        "\n",
        "def get_evalsets():\n",
        "  evalsets = []\n",
        "  if FLAGS.eval_dataset_path:\n",
        "    evalsets.append('%s_dev' % FLAGS.training_dataset)\n",
        "  if FLAGS.guidance_dataset_path and not FLAGS.train_with_guided_parameters:\n",
        "    evalsets.append('%s_guide' % FLAGS.training_dataset)\n",
        "  return evalsets\n",
        "\n",
        "\n",
        "def get_learning_rate_fns():\n",
        "  \"\"\"Get learning rate fns, specified by relevant FLAGS.\n",
        "\n",
        "  Returns:\n",
        "    learning_rate_fn, learning_rate_fn_guided_vars\n",
        "  \"\"\"\n",
        "  warmup_steps = FLAGS.warmup_steps_ratio * FLAGS.num_train_steps\n",
        "  warmup_steps_guided = FLAGS.warmup_steps_ratio_guided_vars * FLAGS.num_train_steps\n",
        "  if FLAGS.warmup_steps != 0:\n",
        "    warmup_steps = FLAGS.warmup_steps\n",
        "    logging.info(\n",
        "        'Overriding FLAGS.warmup_steps_ratio with FLAGS.warmup_steps = %s',\n",
        "        FLAGS.warmup_steps)\n",
        "  learning_rate_fn = utils.create_learning_rate_scheduler(\n",
        "      factors=FLAGS.learning_rate_schedule,\n",
        "      base_learning_rate=FLAGS.learning_rate,\n",
        "      warmup_steps=warmup_steps)\n",
        "  guided_factors = FLAGS.learning_rate_schedule_guided_vars\n",
        "  if FLAGS.use_model_lr_for_guided:\n",
        "    logging.info('Using model learning rate schedule for guided parameters.')\n",
        "    guided_factors = FLAGS.learning_rate_schedule\n",
        "    warmup_steps_guided = warmup_steps\n",
        "  learning_rate_fn_guided_vars = utils.create_learning_rate_scheduler(\n",
        "      factors=guided_factors,\n",
        "      base_learning_rate=FLAGS.learning_rate_guided_vars,\n",
        "      warmup_steps=warmup_steps_guided)\n",
        "\n",
        "  return learning_rate_fn, learning_rate_fn_guided_vars\n",
        "\n",
        "\n",
        "def maybe_reset_guidance_model(\n",
        "    guided_vars_dict,\n",
        "    optimizer_dict,\n",
        "    frequency,\n",
        "    force_reset=False,\n",
        "):\n",
        "  \"\"\"Reset guidance optimizers that are specified to be reset at <frequency>.\n",
        "\n",
        "  Args:\n",
        "    guided_vars_dict: Holds guidance parameter specification.\n",
        "    optimizer_dict: Holds all optimizers\n",
        "    frequency: Which frequency to reset, corresponds to the location of this\n",
        "      method call in the training loop.\n",
        "    force_reset: if True, force the reset even if the given frequency would not\n",
        "\n",
        "  Returns:\n",
        "    guided_vars_dict, optimizer_dict\n",
        "  \"\"\"\n",
        "  if frequency not in ['reset_per_epoch', 'reset_per_batch']:\n",
        "    raise ValueError('Invalid frequency: %s' % frequency)\n",
        "  for var_type in guided_vars_dict:\n",
        "    if guided_vars_dict[var_type][frequency] or force_reset:\n",
        "      if guided_vars_dict[var_type][frequency]:\n",
        "        logging.info('Resetting %s to initialization vals, due to %s tag',\n",
        "                     var_type, frequency)\n",
        "      else:\n",
        "        logging.info(\n",
        "            'Resetting %s to initialization vals, due to reset guidance per steps = %s',\n",
        "            var_type, FLAGS.reset_guidance_per_steps)\n",
        "      guided_vars_dict = guided_parameters.reset_subdict_raw_vars(\n",
        "          guided_vars_dict, var_type)\n",
        "      # import pdb\n",
        "      # pdb.set_trace()\n",
        "      optimizer_dict[guided_parameters.get_opt_key_from_var_type(\n",
        "          var_type)] = jax_utils.replicate(\n",
        "              models.init_optimizer_by_type(\n",
        "                  guided_vars_dict[var_type]['raw_guided_vars'],\n",
        "                  guided_vars_dict[var_type]['optimizer_type'],\n",
        "                  guided_parameters.Granularity.GLOBAL))\n",
        "  return guided_vars_dict, optimizer_dict\n",
        "\n",
        "\n",
        "def output_train_metrics(\n",
        "    t_start,\n",
        "    step,\n",
        "    train_step_metrics,\n",
        "    train_sum_writer,\n",
        "    t_metrics_timer,\n",
        "    guide_step_metrics,\n",
        "    extra_metrics,\n",
        "):\n",
        "  \"\"\"Output training step metrics (may include guide step metrics).\n",
        "\n",
        "\n",
        "  Args:\n",
        "    t_start: For timekeeping\n",
        "    step: Current model step\n",
        "    train_step_metrics: Metrics from the training step.\n",
        "    train_sum_writer: Summary writer for training loop.\n",
        "    t_metrics_timer: keeps track of time since eval\n",
        "    guide_step_metrics: Metrics dict for guide step.\n",
        "    extra_metrics: Extra info regarding training run (epoch, guide ds size, etc)\n",
        "\n",
        "  Returns:\n",
        "    A timestamp, new if metrics were output else t_metrics_timer is passed\n",
        "    through.\n",
        "  \"\"\"\n",
        "  # Periodic metric handling frequency.\n",
        "  logging.info('Output metrics at step: %s', step)\n",
        "  # Training Metrics\n",
        "  train_step_metrics = common_utils.get_metrics(train_step_metrics)\n",
        "\n",
        "  summary = extra_metrics\n",
        "  # Remove already-averaged metrics (learning rate, guidance avgs, etc).\n",
        "  for k in list(train_step_metrics.keys()):\n",
        "    if k.split('dataset_train/')[-1] not in [\n",
        "        'loss', 'accuracy', 'denominator'\n",
        "    ] and k.split('dataset_guide/')[-1] not in [\n",
        "        'loss', 'accuracy', 'denominator'\n",
        "    ]:\n",
        "      # These metrics do not need to be summed / divided, so take them out.\n",
        "      # This is 'learning_rate' and various dppl metrics if present.\n",
        "      try:\n",
        "        summary[k] = train_step_metrics.pop(k).mean()\n",
        "      except AttributeError as e:\n",
        "        logging.warning('Metrics key: %s causes error: %s. Overriding to 0.', k,\n",
        "                        e)\n",
        "        summary[k] = 0.0\n",
        "      pass\n",
        "  # Take the sums across batch for 'loss', 'accuracy', and 'denominator'.\n",
        "  metrics_sums = jax.tree_map(jnp.sum, train_step_metrics)\n",
        "  # Handle weighted loss.\n",
        "  unweighted_keys = [x for x in metrics_sums.keys() if 'training_set_unw/' in x]\n",
        "  if unweighted_keys:\n",
        "    unweighted_metrics_sums = {}\n",
        "    for unw_k in unweighted_keys:\n",
        "      unweighted_metrics_sums[unw_k] = metrics_sums.pop(unw_k)\n",
        "    denominator_unweighted = unweighted_metrics_sums.pop(\n",
        "        'training_set_unw/denominator')\n",
        "    summary.update(\n",
        "        jax.tree_map(lambda x: x / denominator_unweighted,\n",
        "                     unweighted_metrics_sums))  # pylint: disable=cell-var-from-loop\n",
        "  guidance_set_keys = [x for x in metrics_sums.keys() if 'dataset_guide/' in x]\n",
        "  if guidance_set_keys:\n",
        "    unweighted_metrics_sums = {}\n",
        "    for unw_k in guidance_set_keys:\n",
        "      unweighted_metrics_sums[unw_k] = metrics_sums.pop(unw_k)\n",
        "    denominator_unweighted = unweighted_metrics_sums.pop(\n",
        "        'dataset_guide/denominator')\n",
        "    summary.update(\n",
        "        jax.tree_map(lambda x: x / denominator_unweighted,\n",
        "                     unweighted_metrics_sums))  # pylint: disable=cell-var-from-loop\n",
        "\n",
        "  # After this, only loss and accuracy are left in the metrics_sums.\n",
        "  if 'dataset_train/denominator' in metrics_sums:\n",
        "    denominator = metrics_sums.pop('dataset_train/denominator')\n",
        "    summary.update(jax.tree_map(lambda x: x / denominator, metrics_sums))  # pylint: disable=cell-var-from-loop\n",
        "\n",
        "  steps_per_eval = get_train_metrics_freq() if step != 0 else 1\n",
        "  steps_per_sec = steps_per_eval / (time.time() - t_metrics_timer)\n",
        "  if jax.process_index() == 0:\n",
        "    train_sum_writer.scalar('general/steps_per_second', steps_per_sec, step)\n",
        "    for key, val in summary.items():\n",
        "      train_sum_writer.scalar(key, val, step)\n",
        "      if '/loss' in key:\n",
        "        train_sum_writer.scalar('loss/%s' % key, val, step)\n",
        "\n",
        "    train_sum_writer.flush()\n",
        "  train_step_metrics = []\n",
        "  if 'loss' in summary:\n",
        "    utils.log_message(t_start,\n",
        "                      'train in step: %d, loss: %.4f' % (step, summary['loss']))\n",
        "\n",
        "  # Guidance Metrics\n",
        "  if guide_step_metrics:\n",
        "    guide_step_metrics = common_utils.get_metrics(guide_step_metrics)\n",
        "    logging.info('guide_step_metrics: %s', guide_step_metrics)\n",
        "    metrics_summary = {}\n",
        "    for k in list(guide_step_metrics.keys()):\n",
        "      if k.split('/')[-1] not in ['loss', 'accuracy', 'denominator']:\n",
        "        # These metrics do not need to be summed / divided, so take them out.\n",
        "        # This is 'learning_rate' and various guidance metrics if present.\n",
        "        metrics_summary[k] = guide_step_metrics.pop(k).mean()\n",
        "    # Take the sums across batch for 'loss', 'accuracy', and 'denominator'.\n",
        "    metrics_sums = jax.tree_map(jnp.sum, guide_step_metrics)\n",
        "    # After this, only loss and accuracy are left in the metrics_sums.\n",
        "    if 'dataset_guide/denominator' in metrics_sums:\n",
        "      denominator = metrics_sums.pop('dataset_guide/denominator')\n",
        "      # metrics_summary['dataset_guide/denominator'] = denominator\n",
        "      metrics_summary.update(\n",
        "          jax.tree_map(lambda x: x / denominator, metrics_sums))  # pylint: disable=cell-var-from-loop\n",
        "\n",
        "    if jax.process_index() == 0:\n",
        "      train_sum_writer.flush()\n",
        "  guide_step_metrics = []\n",
        "  if 'loss' in summary:\n",
        "    utils.log_message(t_start,\n",
        "                      'train in step: %d, loss: %.4f' % (step, summary['loss']))\n",
        "  return time.time()\n",
        "\n",
        "\n",
        "def get_train_metrics_freq():\n",
        "  metric_frequency = int(FLAGS.num_train_steps /\n",
        "                         min(FLAGS.total_evals, FLAGS.num_train_steps))\n",
        "  if FLAGS.min_eval_freq:\n",
        "    metric_frequency = FLAGS.min_eval_freq\n",
        "  train_metric_freq = max(\n",
        "      FLAGS.max_metric_freq,\n",
        "      metric_frequency // FLAGS.train_metrics_per_eval_metric)\n",
        "  train_metric_freq = min(train_metric_freq, FLAGS.min_train_metric_freq)\n",
        "  return train_metric_freq\n",
        "\n",
        "\n",
        "def output_train_metrics_this_step(step):\n",
        "  train_metric_freq = get_train_metrics_freq()\n",
        "  return step % train_metric_freq == 0\n",
        "\n",
        "\n",
        "def output_eval_metrics_this_step(step):\n",
        "  metric_frequency = int(FLAGS.num_train_steps /\n",
        "                         min(FLAGS.total_evals, FLAGS.num_train_steps))\n",
        "  if FLAGS.min_eval_freq:\n",
        "    metric_frequency = FLAGS.min_eval_freq\n",
        "  return step % metric_frequency == 0\n",
        "\n",
        "\n",
        "def do_eval_all_eval_sets(\n",
        "    step,\n",
        "    p_eval_step,\n",
        "    eval_ds_dict,\n",
        "    eval_sum_writers_dict,\n",
        "    model_optimizer,\n",
        "):\n",
        "  \"\"\"Do eval for all eval datasets.\n",
        "\n",
        "  Args:\n",
        "    step: Current model step.\n",
        "    p_eval_step: Parallelized evaluation step function.\n",
        "    eval_ds_dict: Contains all eval datasets.\n",
        "    eval_sum_writers_dict: Maps eval set names to summary writers.\n",
        "    model_optimizer: The optimizer of the model.\n",
        "\n",
        "  Returns:\n",
        "  \"\"\"\n",
        "  logging.info('Running eval for all evalsets [%s] at step: %s.',\n",
        "               list(eval_sum_writers_dict.keys()), step)\n",
        "  if FLAGS.print_for_colab:\n",
        "    print('Running eval for all evalsets [%s] at step: %s.' %\n",
        "          (list(eval_sum_writers_dict.keys()), step))\n",
        "  for eval_ds_name, eval_summary_writer in eval_sum_writers_dict.items():\n",
        "    logging.info('Starting eval for: %s.', eval_ds_name)\n",
        "\n",
        "    if FLAGS.print_for_colab:\n",
        "      print('Starting eval for: %s.' % eval_ds_name)\n",
        "    eval_ds = eval_ds_dict[eval_ds_name]\n",
        "    t_before_eval = time.time()\n",
        "    eval_summary = do_eval_single(step, eval_ds, eval_ds_name, p_eval_step,\n",
        "                                  model_optimizer)\n",
        "    logging.info('Finished eval for: %s.', eval_ds_name)\n",
        "\n",
        "    eval_duration = time.time() - t_before_eval\n",
        "    if FLAGS.print_for_colab:\n",
        "      print('finished eval for %s in duration %s' %\n",
        "            (eval_ds_name, eval_duration))\n",
        "    if jax.process_index() == 0:\n",
        "      eval_summary_writer.scalar('dataset_eval/duration', eval_duration, step)\n",
        "      for key, val in eval_summary.items():\n",
        "        # If we slipped the guidance dataset into the evals (we are training\n",
        "        # without guided parameters), make sure we pull it back out into it's\n",
        "        # correct category for comparison to guided runs.\n",
        "        # eval_ds_simple_name = eval_ds_name.split('-')[-1]\n",
        "        if 'guide-' in eval_ds_name:\n",
        "          key = key.replace('dataset_eval', 'dataset_guide')\n",
        "        eval_summary_writer.scalar(key, val, step)\n",
        "        if '/loss' in key:\n",
        "          eval_summary_writer.scalar('loss/%s' % eval_ds_name.split('-')[-1],\n",
        "                                     val, step)\n",
        "\n",
        "      eval_summary_writer.flush()\n",
        "  logging.info('Finished eval of all sets at step %s', step)\n",
        "\n",
        "\n",
        "def jax_synchronize_hosts():\n",
        "  \"\"\"Ensure all jax hosts synchronize before exiting.\"\"\"\n",
        "  if jax.process_count() > 1:\n",
        "    # Make sure all hosts stay up until the end of main.\n",
        "    x = jnp.ones([jax.local_device_count()])\n",
        "    x = jax.device_get(jax.pmap(lambda x: jax.lax.psum(x, 'i'), 'i')(x))\n",
        "    assert x[0] == jax.device_count()\n"
      ],
      "metadata": {
        "id": "y5hsPRxRWRaw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "HWMzDfXNj1vB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup Training\n",
        "\n",
        "Flags shared by all training runs. Modify as needed."
      ],
      "metadata": {
        "id": "6XAGrqTZj2Uy"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "spm_model_path = '/content/data/%s.32k.spm.model' % LANG_PAIR\n",
        "train_data_path_tfr='/content/data/%s_train_99percent.tfr-00000-of-00001' % LANG_PAIR\n",
        "guide_data_path_tfr='/content/data/%s_guide_1percent.tfr-00000-of-00001' % LANG_PAIR\n",
        "dev_data_path_tfr='/content/data/%s_dev.tfr-00000-of-00001' % LANG_PAIR\n",
        "\n",
        "FLAGS = flags.FLAGS"
      ],
      "metadata": {
        "id": "MMTUwQEUX804"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def set_colab_default_flags(size=None):\n",
        "  \n",
        "  FLAGS.print_for_colab=True\n",
        "\n",
        "  # local paths\n",
        "  FLAGS.vocab_path=spm_model_path\n",
        "  FLAGS.training_dataset_path = train_data_path_tfr\n",
        "  FLAGS.guidance_dataset_path = guide_data_path_tfr\n",
        "  FLAGS.eval_dataset_path = dev_data_path_tfr\n",
        "\n",
        "  # model\n",
        "  FLAGS.dropout_rate = 0.1\n",
        "  FLAGS.attention_dropout_rate = 0.1\n",
        "\n",
        "  # eval / data\n",
        "  FLAGS.training_dataset = LANG_PAIR\n",
        "  FLAGS.batch_size = 16\n",
        "  FLAGS.eval_batch_size = 16\n",
        "  FLAGS.total_evals = 5\n",
        "  FLAGS.testsets=[]\n",
        "  FLAGS.warmup_steps=1000\n",
        "  FLAGS.min_eval_freq = 250\n",
        "  FLAGS.do_eval=True\n",
        "  FLAGS.guide_batch_size = 16 \n",
        "  FLAGS.guide_batch_update_freq = 'NEVER'\n",
        "\n",
        "\n",
        "  # schedules\n",
        "  FLAGS.learning_rate_schedule = 'constant * linear_warmup * rsqrt_decay'\n",
        "  FLAGS.learning_rate_schedule_guided_vars = 'constant'\n",
        "  \n",
        "\n",
        "\n",
        "  # ckpts\n",
        "  FLAGS.save_checkpoints = False\n",
        "  FLAGS.guided_ckpt_min_freq = 1000\n",
        "  FLAGS.model_ckpt_min_freq=500\n",
        "  FLAGS.max_train_epochs = 10000\n",
        "\n",
        "  # inference\n",
        "  FLAGS.do_predict=False\n",
        "\n",
        "  #training\n",
        "  FLAGS.num_train_steps = 2000\n",
        "  FLAGS.learning_rate = 0.4\n",
        "\n",
        "  # guidance\n",
        "  FLAGS.train_with_guided_parameters = True\n",
        "  FLAGS.learning_rate_guided_vars =  0.003\n",
        "  FLAGS.guided_params_optimizer_type='adam'\n",
        "  FLAGS.guided_hparam_types='all'\n",
        "  FLAGS.model_optimizer_type='lamb'\n",
        "\n",
        "  if size=='tiny':\n",
        "    FLAGS.emb_dim=64\n",
        "    FLAGS.num_layers=2\n",
        "    FLAGS.qkv_dim=64\n",
        "    FLAGS.mlp_dim=128\n",
        "  elif size=='small':\n",
        "    FLAGS.emb_dim=96\n",
        "    FLAGS.num_layers=3\n",
        "    FLAGS.qkv_dim=64\n",
        "    FLAGS.mlp_dim=256\n",
        "  elif size=='medium':\n",
        "    FLAGS.emb_dim=384\n",
        "    FLAGS.num_layers=4\n",
        "    FLAGS.qkv_dim=256\n",
        "    FLAGS.mlp_dim=1024\n",
        "  elif size=='large' or size=='paper':\n",
        "    FLAGS.emb_dim=1536\n",
        "    FLAGS.num_layers=6\n",
        "    FLAGS.qkv_dim=1024\n",
        "    FLAGS.mlp_dim=4096"
      ],
      "metadata": {
        "id": "eeV0iBARPAvT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def training_main():\n",
        "  FLAGS(sys.argv[1:])\n",
        "  maybe_remove_finished_training(FLAGS.model_dir)\n",
        "  t_start = time.time()\n",
        "\n",
        "  # # Necessary to make it work on GPU without OOM.\n",
        "  # # https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html#common-causes-of-oom-failures\n",
        "  # tf.config.experimental.set_visible_devices([], 'GPU')\n",
        "\n",
        "  # Number of local devices for this host.\n",
        "  n_devices = jax.local_device_count()\n",
        "  logging.info('JAX host: process_index %d / process_count %d',\n",
        "                jax.process_index(), jax.process_count())\n",
        "\n",
        "  if FLAGS.batch_size % n_devices:\n",
        "    raise ValueError(\n",
        "        'Batch size (%s) must be divisible by the number of devices (%s)' %\n",
        "        (FLAGS.batch_size, n_devices))\n",
        "\n",
        "  if FLAGS.jax_backend_target:\n",
        "    jax.config.FLAGS.jax_xla_backend = 'tpu_driver'\n",
        "    jax.config.FLAGS.jax_backend_target = FLAGS.jax_backend_target\n",
        "\n",
        "  # Output functions (summary writers, writing out flags file) need only be\n",
        "  # done once.\n",
        "  train_sum_writer, eval_sum_writers_dict = set_up_summary_writers()\n",
        "\n",
        "  # Load data-parameters dict.\n",
        "  utils.log_message(t_start, 'Maybe loading guided_vars_dict.')\n",
        "\n",
        "  optimizer_dict = {}\n",
        "  model_optimizer, current_step, rng, model = restore_or_init_model(t_start)\n",
        "  optimizer_dict['model'] = jax_utils.replicate(model_optimizer)\n",
        "  # Replicate optimizer.\n",
        "  del model_optimizer  # Only use through dict.\n",
        "  del model  # Only reference through dict.\n",
        "\n",
        "  extra_metrics = {\n",
        "      'general/num_params_model':\n",
        "          sum(x.size for x in jax.tree_leaves(optimizer_dict['model'].target)),\n",
        "  }\n",
        "\n",
        "  # this allows us to set separate learning_rate_scalar values for each guided\n",
        "  # parameter via FLAGS. All guided params share a learning rate, but this\n",
        "  # scalar is multiplied by that learning rate for each one separately, allowing\n",
        "  # the params to take different learning rates.\n",
        "  learning_rate_scalar_override = {\n",
        "      'decay_rate': 1,\n",
        "      'beta1': FLAGS.ghp_lrs_beta1,\n",
        "      'eps': FLAGS.ghp_lrs_eps,\n",
        "      'label_smoothing': FLAGS.ghp_lrs_ls,\n",
        "      'learning_rate': FLAGS.ghp_lrs_lr,\n",
        "      'weight_decay': FLAGS.ghp_lrs_wd,\n",
        "  }\n",
        "\n",
        "  guided_vars_dict = None\n",
        "\n",
        "  if FLAGS.train_with_guided_parameters:\n",
        "    opt_kwargs = _get_default_opt_kwargs(\n",
        "        model_optimizer_type=FLAGS.model_optimizer_type,\n",
        "        lr=1,\n",
        "    )\n",
        "    guided_vars_dict = guided_parameters.get_guided_vars_dict(\n",
        "        guided_hparam_types=FLAGS.guided_hparam_types,\n",
        "        model_opt_type=FLAGS.model_optimizer_type,\n",
        "        guided_opt_type=FLAGS.guided_params_optimizer_type,\n",
        "        init_dict=opt_kwargs,\n",
        "        learning_rate_scalar_override=learning_rate_scalar_override,\n",
        "    )\n",
        "\n",
        "    if jax.process_index() == 0:\n",
        "      guided_parameters.save_guided_vars_dict(guided_vars_dict, FLAGS.model_dir)\n",
        "  # Write flags/guided_vars_dict state to CNS.\n",
        "  write_flags_files(FLAGS.model_dir, guided_vars_dict)\n",
        "  utils.log_message(\n",
        "      t_start,\n",
        "      'Loaded model, current training step: %s, model optimizer step (total steps since random init): %s'\n",
        "      % (current_step, get_step_from_opt(optimizer_dict['model'])))\n",
        "\n",
        "  print('NUM_PARAMS:%s' % sum(x.size for x in jax.tree_leaves(optimizer_dict['model'].target)))\n",
        "  # raise ValueError\n",
        "  # Set learning rate fns.\n",
        "  learning_rate_fn, learning_rate_fn_hyperparams = get_learning_rate_fns()\n",
        "\n",
        "  # Load Dataset\n",
        "  utils.log_message(t_start, 'Initializing dataset.')\n",
        "  if FLAGS.print_for_colab:\n",
        "    print('Initializing dataset.')\n",
        "  if FLAGS.training_dataset_path:\n",
        "    logging.info('FLAGS.training_dataset_path  %s', FLAGS.training_dataset_path)\n",
        "\n",
        "    train_ds_kwargs = {\n",
        "        'batch_size': FLAGS.batch_size,\n",
        "        'deterministic': FLAGS.print_for_colab,\n",
        "        'max_length': FLAGS.max_target_length,\n",
        "        'pack': True,\n",
        "        'random_seed': FLAGS.jax_random_seed,\n",
        "        'repeat': True,\n",
        "        'vocab_path': FLAGS.vocab_path,\n",
        "    }\n",
        "    train_ds = data.get_prepacked_examples(\n",
        "        file_pattern=FLAGS.training_dataset_path, **train_ds_kwargs)\n",
        "  if FLAGS.guidance_dataset_path:\n",
        "    logging.info('FLAGS.guidance_dataset_path  %s', FLAGS.guidance_dataset_path)\n",
        "\n",
        "    eval_ds_dict = {}\n",
        "    guide_ds_kwargs = {\n",
        "        'batch_size': FLAGS.guide_batch_size,\n",
        "        'deterministic': FLAGS.print_for_colab,\n",
        "        'drop_remainder': False,\n",
        "        'max_length': FLAGS.max_target_length,\n",
        "        'pack': True,\n",
        "        'random_seed': FLAGS.jax_random_seed,\n",
        "        'repeat': False,\n",
        "        'vocab_path': FLAGS.vocab_path,\n",
        "        'shard_data': False,\n",
        "    }\n",
        "    guide_ds = data.get_prepacked_examples(\n",
        "        file_pattern=FLAGS.guidance_dataset_path, **guide_ds_kwargs)\n",
        "    if not FLAGS.train_with_guided_parameters:\n",
        "      eval_ds_dict['eval-%s_guide' % FLAGS.training_dataset] = guide_ds\n",
        "    if FLAGS.guide_batch_update_freq == 'NEVER' and FLAGS.training_dataset != 'guide':\n",
        "      logging.info('Truncating guidance dataset to a single batch')\n",
        "      if FLAGS.print_for_colab:\n",
        "        print('Truncating guidance dataset to a single batch')\n",
        "      guide_ds = guide_ds.take(jax.local_device_count()).repeat()\n",
        "  if FLAGS.eval_dataset_path:\n",
        "    logging.info('FLAGS.eval_dataset_path  %s', FLAGS.eval_dataset_path)\n",
        "    eval_ds_kwargs = {\n",
        "        'batch_size': FLAGS.eval_batch_size,\n",
        "        'deterministic': True,\n",
        "        'drop_remainder': False,\n",
        "        'max_length': FLAGS.max_target_length,\n",
        "        'pack': True,\n",
        "        'random_seed': FLAGS.jax_random_seed,\n",
        "        'repeat': False,\n",
        "        'shard_data': False,\n",
        "        'vocab_path': FLAGS.vocab_path,\n",
        "    }\n",
        "    eval_ds = data.get_prepacked_examples(\n",
        "        file_pattern=FLAGS.eval_dataset_path,\n",
        "        **eval_ds_kwargs).take(jax.local_device_count())\n",
        "    eval_ds_dict['eval-%s_dev' % FLAGS.training_dataset] = eval_ds\n",
        "  logging.info('train keys: %s')\n",
        "  if FLAGS.print_for_colab:\n",
        "    print('train keys:')\n",
        "    for x in next(iter(train_ds)).keys():\n",
        "      print('\\t', x)\n",
        "    if FLAGS.train_with_guided_parameters:\n",
        "      print('guide keys: %s')\n",
        "      for x in next(iter(guide_ds)).keys():\n",
        "        print('\\t', x)\n",
        "\n",
        "  guide_batch_unq_ex = 0\n",
        "  train_batch_unq_ex = 0\n",
        "  if FLAGS.train_with_guided_parameters:\n",
        "    guide_iter = iter(guide_ds)  # This is an infinite iterator.\n",
        "    guide_batch = None\n",
        "    if FLAGS.guide_batch_update_freq == 'NEVER':\n",
        "      guide_batch = batch_to_numpy(next(guide_iter))\n",
        "      # print guide_batch unique examples count\n",
        "      guide_batch_unq_ex = data.get_unique_examples(guide_batch)\n",
        "      logging.info('Guide batch unique examples: %s', guide_batch_unq_ex)\n",
        "      if FLAGS.print_for_colab:\n",
        "        print('Guide batch unique examples:', guide_batch_unq_ex)\n",
        "      guide_batch = common_utils.shard(guide_batch)\n",
        "\n",
        "    if FLAGS.print_for_colab:\n",
        "      print('Guide batch update freq:', FLAGS.guide_batch_update_freq)\n",
        "\n",
        "  # Load guided params state.\n",
        "  if FLAGS.train_with_guided_parameters:\n",
        "    guided_optimizers = restore_or_init_guided_optimizers(\n",
        "        t_start, guided_vars_dict)\n",
        "    optimizer_dict.update(guided_optimizers)\n",
        "\n",
        "  # set up model pmapped update fns\n",
        "  if not FLAGS.train_with_guided_parameters:\n",
        "    p_update_model_step = jax.pmap(\n",
        "        functools.partial(\n",
        "            update_model_step,\n",
        "            learning_rate_fn=learning_rate_fn,\n",
        "            use_bfloat16=FLAGS.use_bfloat16),\n",
        "        axis_name='batch')\n",
        "  # Set up guided vars pmapped step fns.\n",
        "  elif FLAGS.train_with_guided_parameters:\n",
        "    utils.log_message(t_start, 'Setting up p_update_model_and_guidance_step.')\n",
        "    p_update_model_and_guidance_step = jax.pmap(\n",
        "        functools.partial(\n",
        "            update_model_and_guidance_step,\n",
        "            learning_rate_fn=learning_rate_fn,\n",
        "            learning_rate_fn_guided_vars=learning_rate_fn_hyperparams,\n",
        "            guided_vars_dict=guided_vars_dict,\n",
        "            grad_clip_limit=FLAGS.grad_clip_limit,\n",
        "            use_bfloat16=FLAGS.use_bfloat16,\n",
        "        ),\n",
        "        axis_name='batch')\n",
        "\n",
        "  p_eval_step = jax.pmap(\n",
        "      functools.partial(eval_step, use_bfloat16=FLAGS.use_bfloat16),\n",
        "      axis_name='batch')\n",
        "  # We init the first set of dropout PRNG keys, but update it afterwards inside\n",
        "  # the main pmap'd training update for performance.\n",
        "  dropout_rngs = random.split(rng, n_devices)\n",
        "\n",
        "  utils.log_message(t_start, 'Starting training loop.')\n",
        "  if FLAGS.print_for_colab:\n",
        "    print(t_start, 'Starting training loop.')\n",
        "  train_step_metrics = []\n",
        "  guide_step_metrics = []\n",
        "  metrics_train = None\n",
        "  metrics_joint = None\n",
        "  metrics_guide = None\n",
        "\n",
        "  t_metrics_timer = time.time()\n",
        "  # Training loop\n",
        "  delta_guide_loss_dict = {}\n",
        "  steps_since_last_model_ckpt = 0\n",
        "  steps_since_last_dp_ckpt = 0\n",
        "  steps_since_preemption = 0\n",
        "  current_epoch = guided_parameters.load_epoch(FLAGS.model_dir)\n",
        "  logging.info('Loaded epoch: %s', current_epoch)\n",
        "\n",
        "  if current_step == 0 and FLAGS.save_checkpoint_at_init:\n",
        "    # This is not a restart after pre-emption.\n",
        "    logging.info('saving checkpoint 0: current_step == %s', current_step)\n",
        "    save_all_checkpoints(t_start, optimizer_dict, current_step)\n",
        "    # If not random-init, do eval before starting training.\n",
        "    if FLAGS.init_checkpoint and FLAGS.do_eval:\n",
        "      logging.info('Doing eval before training.')\n",
        "      do_eval_all_eval_sets(current_step, p_eval_step, eval_ds_dict,\n",
        "                            eval_sum_writers_dict, optimizer_dict['model'])\n",
        "\n",
        "  # Loop through epochs.\n",
        "  for epoch in range(current_epoch, FLAGS.max_train_epochs):\n",
        "    logging.info('Starting epoch loop for epoch %s', epoch)\n",
        "    if FLAGS.print_for_colab:\n",
        "      print('Starting epoch loop for epoch %s' % epoch)\n",
        "    if jax.process_index() == 0:\n",
        "      guided_parameters.save_epoch(epoch, FLAGS.model_dir)\n",
        "    if FLAGS.num_train_steps > 0 and current_step > FLAGS.num_train_steps:\n",
        "      logging.info(\n",
        "          'Breaking out of trianing loop as current step %s > FLAGS.num_train_steps %s',\n",
        "          current_step, FLAGS.num_train_steps)\n",
        "      if FLAGS.print_for_colab:\n",
        "        print(\n",
        "            f'Breaking out of trianing loop as current step {current_step} > FLAGS.num_train_steps {FLAGS.num_train_steps}'\n",
        "        )\n",
        "      break\n",
        "\n",
        "    steps_per_epoch = 0\n",
        "    # Refresh the non-repeating (single epoch) train_ds iterator.\n",
        "    if FLAGS.train_with_guided_parameters:\n",
        "      guide_batch, guide_iter, guide_batch_unq_ex = maybe_get_new_guide_batch(\n",
        "          guide_batch, guide_iter, 'PER_EPOCH', guide_batch_unq_ex)\n",
        "\n",
        "      # maybe_reset_guidance_model_epoch\n",
        "      # this must go after the filtering! otherwise no filtering will occur due\n",
        "      # to reset values. Don't filter on the first (0th) epoch.\n",
        "      if epoch != 0:\n",
        "        guided_vars_dict, optimizer_dict = maybe_reset_guidance_model(\n",
        "            guided_vars_dict, optimizer_dict, 'reset_per_epoch')\n",
        "    train_iter = iter(train_ds)\n",
        "\n",
        "    # Per epoch loop.\n",
        "    epoch_step = 0\n",
        "    finish_epoch = False\n",
        "    for train_batch in train_iter:\n",
        "      if FLAGS.num_train_steps > 0 and current_step > FLAGS.num_train_steps:\n",
        "        break\n",
        "      if finish_epoch:\n",
        "        finish_epoch = False\n",
        "        break\n",
        "      # Shard data for multiple devices.\n",
        "      train_batch = batch_to_numpy(train_batch)\n",
        "      if FLAGS.guide_with_train_loss:\n",
        "        guide_batch = common_utils.shard(train_batch)\n",
        "      train_batch_unq_ex = data.get_unique_examples(train_batch)\n",
        "\n",
        "      epoch_step += 1\n",
        "      if FLAGS.num_train_steps > 0 and current_step > FLAGS.num_train_steps:\n",
        "        logging.info('num_train_steps (%s) reached', FLAGS.num_train_steps)\n",
        "        save_all_checkpoints(t_start, optimizer_dict, current_step)\n",
        "        logging.info('Saving ckpts and quitting.')\n",
        "        break\n",
        "      current_step += 1\n",
        "      steps_per_epoch += 1\n",
        "      steps_since_preemption += 1\n",
        "      if current_step % 1000 == 0:\n",
        "        utils.log_message(t_start, 'Current_step: %s' % current_step)\n",
        "\n",
        "      if FLAGS.train_with_guided_parameters:\n",
        "        if FLAGS.reset_guidance_per_steps and current_step % FLAGS.reset_guidance_per_steps == 0:\n",
        "          guided_vars_dict, optimizer_dict = maybe_reset_guidance_model(\n",
        "              guided_vars_dict,\n",
        "              optimizer_dict,\n",
        "              'reset_per_epoch',\n",
        "              force_reset=True)\n",
        "          if FLAGS.print_for_colab:\n",
        "            print('resetting guided vars at step: %s' % current_step)\n",
        "        guide_batch, guide_iter, guide_batch_unq_ex = maybe_get_new_guide_batch(\n",
        "            guide_batch, guide_iter, 'PER_TRAINING_STEP', guide_batch_unq_ex)\n",
        "\n",
        "        if current_step % (max(FLAGS.num_train_steps // 10,\n",
        "                                1)) == 0 or current_step < 10:\n",
        "          logging.info('train step %s', current_step)\n",
        "          if FLAGS.print_for_colab:\n",
        "            print('train step %s' % current_step)\n",
        "        sharded_train_batch = common_utils.shard(train_batch)\n",
        "\n",
        "        # Either joint optimization, or run validation and train steps in seq.\n",
        "        joint_step_kwargs = {\n",
        "            'train_batch': sharded_train_batch,\n",
        "            'dropout_rng': dropout_rngs,\n",
        "            'guide_batch': guide_batch,\n",
        "        }\n",
        "        if output_train_metrics_this_step(current_step):\n",
        "          joint_step_kwargs.update({\n",
        "              'output_metrics': jnp.array([True] * jax.local_device_count()),\n",
        "          })\n",
        "        optimizer_dict, metrics_joint, dropout_rngs = p_update_model_and_guidance_step(\n",
        "            optimizer_dict, **joint_step_kwargs)\n",
        "      else:\n",
        "        if FLAGS.train_with_guided_parameters:\n",
        "          # using update_model_multiweighted_train_loss_dict()\n",
        "          optimizer_dict, metrics_train, dropout_rngs = p_update_model_step(\n",
        "              optimizer_dict,\n",
        "              train_batch=common_utils.shard(train_batch),\n",
        "              dropout_rng=dropout_rngs)\n",
        "        else:\n",
        "          # using update_model_step()\n",
        "          optimizer_dict, metrics_train, dropout_rngs = p_update_model_step(\n",
        "              optimizer_dict,\n",
        "              train_batch=common_utils.shard(train_batch),\n",
        "              dropout_rng=dropout_rngs)\n",
        "\n",
        "      if metrics_train:\n",
        "        train_step_metrics.append(metrics_train)\n",
        "      if metrics_guide:\n",
        "        guide_step_metrics.append(metrics_guide)\n",
        "      if metrics_joint:\n",
        "        train_step_metrics.append(metrics_joint)\n",
        "\n",
        "      extra_metrics.update({\n",
        "          'general/steps_since_preemption': steps_since_preemption,\n",
        "          'general/epoch': epoch,\n",
        "          'dataset_train/train_batch_unq_ex': train_batch_unq_ex,\n",
        "          'dataset_guide/guide_batch_unq_ex': guide_batch_unq_ex,\n",
        "          'general/steps_per_epoch': steps_per_epoch,\n",
        "          'hparams/dropout_rate': FLAGS.dropout_rate,\n",
        "          'hparams/attention_dropout_rate': FLAGS.attention_dropout_rate,\n",
        "          'hparams/label_smoothing': FLAGS.label_smoothing,\n",
        "      })\n",
        "      # TODO(lichtarge): Calculate change in guidance loss, add to metrics.\n",
        "      for key, v in delta_guide_loss_dict.items():\n",
        "        extra_metrics.update({'delta_guide_loss-%s' % key: v})\n",
        "      if output_train_metrics_this_step(current_step):\n",
        "        logging.info('Output train metrics on step: %s', current_step)\n",
        "        t_metrics_timer = output_train_metrics(\n",
        "            t_start,\n",
        "            current_step,\n",
        "            train_step_metrics,\n",
        "            train_sum_writer,\n",
        "            t_metrics_timer,\n",
        "            guide_step_metrics,\n",
        "            extra_metrics,\n",
        "        )\n",
        "        extra_metrics = {}\n",
        "      if output_eval_metrics_this_step(current_step):\n",
        "        if FLAGS.do_eval:\n",
        "          do_eval_all_eval_sets(current_step, p_eval_step, eval_ds_dict,\n",
        "                                eval_sum_writers_dict, optimizer_dict['model'])\n",
        "      train_step_metrics, guide_step_metrics = [], []\n",
        "      steps_since_last_model_ckpt, steps_since_last_dp_ckpt = maybe_save_ckpts(\n",
        "          t_start, optimizer_dict, current_step, steps_since_last_model_ckpt,\n",
        "          steps_since_last_dp_ckpt)\n",
        "\n",
        "    logging.info('Epoch finished! current_step: %s, epoch %s, steps per %s',\n",
        "                  current_step, epoch, steps_per_epoch)\n",
        "    if FLAGS.print_for_colab:\n",
        "      print('EPOCH! current_step: %s, epoch %s, steps per %s' %\n",
        "            (current_step, epoch, steps_per_epoch))\n",
        "    if FLAGS.save_ckpt_per_epoch:\n",
        "      logging.info('saving all ckpts at end of epoch')\n",
        "      save_all_checkpoints(t_start, optimizer_dict, current_step)\n",
        "      steps_since_last_model_ckpt = 0\n",
        "\n",
        "  # Done with training loop.\n",
        "  write_finished_training(FLAGS.model_dir)\n",
        "\n",
        "  if FLAGS.print_for_colab:\n",
        "    print('Finished training loop.')\n",
        "  jax_synchronize_hosts()\n",
        "\n"
      ],
      "metadata": {
        "id": "kUjnopHIw2ya"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run Training\n",
        "\n",
        "Demonstrates training with guided hyper-parameters for learning rate, momentum, and the combination of the two. Uses a tiny model to avoid OOM. Each run takes approx 10-20 min on GPU.\n",
        "\n",
        "Runs the following:\n",
        "\n",
        "\n",
        "1.   Baseline (no guided variables)\n",
        "2.   Guided learning rate\n",
        "3.   Guided momentum\n",
        "4.   Guided learning rate and momentum\n",
        "\n",
        "Note that if the runtime is under memory pressure, the tensorboards from earlier runs may die. To open them again once all training is done, rerun those tensorboards with the FLAGS.model_dir arg replaced with the dir of the specific model you are interested in.\n",
        "\n"
      ],
      "metadata": {
        "id": "8cPSpYkdDKQB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        ""
      ],
      "metadata": {
        "id": "ruVmN4wQc1jA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Unguided baseline"
      ],
      "metadata": {
        "id": "KvzLMDosc2vw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# This takes about 12 mins on GPU.\n",
        "\n",
        "# Sets model_dir to new name per run.\n",
        "set_colab_default_flags(size='tiny')\n",
        "FLAGS.model_dir='../content/models_baseline_%s_baseline__%s' % (LANG_PAIR, time.strftime('%Y_%m_%d_%H%M'))\n",
        "\n",
        "# Set the guided hyper-parameter and meta-learning rate.\n",
        "FLAGS.guided_hparam_types='beta1'\n",
        "FLAGS.learning_rate_guided_vars = 0.003\n",
        "\n",
        "# Set number of steps.\n",
        "FLAGS.num_train_steps=10000\n",
        "\n",
        "FLAGS.train_with_guided_parameters=0\n",
        "\n",
        "training_main()\n",
        "print(FLAGS.model_dir)\n",
        "\n"
      ],
      "metadata": {
        "id": "u0kwbgagLrcf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(FLAGS.model_dir)\n",
        "%tensorboard --logdir={FLAGS.model_dir} --port=0"
      ],
      "metadata": {
        "id": "HOZLQivTL9C_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Guide LR"
      ],
      "metadata": {
        "id": "oooNafwKc6T8"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# This takes about 20 mins on GPU.\n",
        "\n",
        "set_colab_default_flags(size='tiny')\n",
        "\n",
        "# Sets model_dir to new name per run.\n",
        "FLAGS.model_dir='../content/models_baseline_%s_lr__%s' % (LANG_PAIR, time.strftime('%Y_%m_%d_%H%M'))\n",
        "\n",
        "# Set the guided hyper-parameter and meta-learning rate.\n",
        "FLAGS.guided_hparam_types='learning_rate'\n",
        "FLAGS.learning_rate_guided_vars = 0.003\n",
        "\n",
        "# Set number of steps.\n",
        "FLAGS.num_train_steps=10000\n",
        "\n",
        "FLAGS.train_with_guided_parameters=1\n",
        "\n",
        "training_main()\n",
        "print(FLAGS.model_dir)\n",
        "\n"
      ],
      "metadata": {
        "id": "uFjXTRAY-IM9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Guide momentum"
      ],
      "metadata": {
        "id": "8LyhRCmJdBth"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "print(FLAGS.model_dir)\n",
        "%tensorboard --logdir={FLAGS.model_dir} --port=0\n"
      ],
      "metadata": {
        "id": "9dTH6BZGM8eq"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# This takes about 20 mins on GPU.\n",
        "\n",
        "set_colab_default_flags(size='tiny')\n",
        "\n",
        "# Sets model_dir to new name per run.\n",
        "FLAGS.model_dir='../content/models_%s_b1__%s' % (LANG_PAIR, time.strftime('%Y_%m_%d_%H%M'))\n",
        "\n",
        "# Set the guided hyper-parameter and meta-learning rate.\n",
        "FLAGS.guided_hparam_types='beta1'\n",
        "FLAGS.learning_rate_guided_vars = 0.003\n",
        "\n",
        "# Set number of steps.\n",
        "FLAGS.num_train_steps=10000\n",
        "\n",
        "training_main()\n"
      ],
      "metadata": {
        "id": "-2iRCKh_dlTL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%tensorboard --logdir={FLAGS.model_dir} --port=0\n"
      ],
      "metadata": {
        "id": "ywhkfwjvdlQD"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Guide momentum and LR"
      ],
      "metadata": {
        "id": "VzZw6hswc9zJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# This takes about 20 mins on GPU.\n",
        "\n",
        "set_colab_default_flags(size='tiny')\n",
        "\n",
        "# Sets model_dir to new name per run.\n",
        "FLAGS.model_dir='../content/models_%s_lrb1__%s' % (LANG_PAIR, time.strftime('%Y_%m_%d_%H%M'))\n",
        "\n",
        "# Set the guided hyper-parameter and meta-learning rate.\n",
        "FLAGS.guided_hparam_types='beta1,learning_rate'\n",
        "FLAGS.learning_rate_guided_vars = 0.003\n",
        "FLAGS.train_with_guided_parameters=1\n",
        "\n",
        "# Set number of steps.\n",
        "FLAGS.num_train_steps=10000\n",
        "\n",
        "training_main()\n"
      ],
      "metadata": {
        "id": "9Bf6kfKv6vrY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%tensorboard --logdir={FLAGS.model_dir} --port=0\n"
      ],
      "metadata": {
        "id": "Sc8H4YGG6wrs"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "-0esjq3luJYi",
        "JndnmDMp66FL",
        "dbHoE4sNB2h7",
        "Hd0GYpo5jsD4",
        "6XAGrqTZj2Uy"
      ],
      "name": "Github: gradient based tuning demo",
      "provenance": [],
      "toc_visible": true,
      "history_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}